diff --git a/.release-please-manifest.json b/.release-please-manifest.json index f1515fa764..6ed83d4523 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "1.80.0" + ".": "1.81.0" } diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b1887b781..87efbf88f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,37 @@ # Changelog +## [1.81.0](https://github.com/googleapis/python-aiplatform/compare/v1.80.0...v1.81.0) (2025-02-18) + + +### Features + +* A new field `create_time` is added to message `.google.cloud.aiplatform.v1.GenerateContentResponse` ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A new field `response_id` is added to message `.google.cloud.aiplatform.v1.GenerateContentResponse` ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A new field `unversioned_package_disabled` is added to message `.google.api.PythonSettings` ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Add additional Probe options to v1 model.proto ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Add Notebooks Runtime Software Configuration ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Add REST Interceptors which support reading metadata ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Add RolloutOptions to DeployedModel in v1beta1 endpoint.proto, add additional Probe options in v1beta1 model.proto ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Add support for reading selective GAPIC generation methods from service YAML ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* EvaluateDataset API v1beta1 initial release ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Use default file name if one isn't provided when uploading evaluation results to GCS. ([d86d87a](https://github.com/googleapis/python-aiplatform/commit/d86d87a87ed6b22f39b95f7b09365570905f017a)) + + +### Bug Fixes + +* **deps:** Require grpc-google-iam-v1>=0.14.0 ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* Prediction stream raw predict url for non dedicated endpoint ([b91a8aa](https://github.com/googleapis/python-aiplatform/commit/b91a8aa95e89f93150522672b5ee367c8f3211ac)) + + +### Documentation + +* A comment for field `content` in message `.google.api.Page` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A comment for field `filter` in message `.google.cloud.aiplatform.v1.ListNotebookRuntimesRequest` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A comment for field `filter` in message `.google.cloud.aiplatform.v1.ListNotebookRuntimeTemplatesRequest` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A comment for field `filter` in message `.google.cloud.aiplatform.v1beta1.ListNotebookRuntimesRequest` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A comment for field `filter` in message `.google.cloud.aiplatform.v1beta1.ListNotebookRuntimeTemplatesRequest` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) +* A comment for message `RoutingRule` is changed ([b176d13](https://github.com/googleapis/python-aiplatform/commit/b176d13092bba744417a43d944f3c7639f73d5f6)) + ## [1.80.0](https://github.com/googleapis/python-aiplatform/compare/v1.79.0...v1.80.0) (2025-02-11) diff --git a/google/cloud/aiplatform/gapic_version.py b/google/cloud/aiplatform/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/gapic_version.py +++ b/google/cloud/aiplatform/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index 8b13f95502..e46915bf28 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -2583,6 +2583,7 @@ def stream_raw_predict( timeout: Optional[float] = None, ) -> Iterator[requests.models.Response]: """Makes a streaming prediction request using arbitrary headers. + For custom model, this method is only supported for dedicated endpoint. Example usage: ``` @@ -2632,7 +2633,7 @@ def stream_raw_predict( if self.stream_raw_predict_request_url is None: self.stream_raw_predict_request_url = f"https://{self.location}-{constants.base.API_BASE_PATH}/v1/projects/{self.project}/locations/{self.location}/endpoints/{self.name}:streamRawPredict" - url = self.raw_predict_request_url + url = self.stream_raw_predict_request_url if use_dedicated_endpoint: self._sync_gca_resource_if_skipped() diff --git a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py index 1536b37754..19ae32e0ad 100644 --- a/google/cloud/aiplatform/version.py +++ b/google/cloud/aiplatform/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.80.0" +__version__ = "1.81.0" diff --git a/google/cloud/aiplatform_v1/__init__.py b/google/cloud/aiplatform_v1/__init__.py index 0740781169..b397fcba8d 100644 --- a/google/cloud/aiplatform_v1/__init__.py +++ b/google/cloud/aiplatform_v1/__init__.py @@ -703,9 +703,12 @@ from .types.model_service import ListModelEvaluationsResponse from .types.model_service import ListModelsRequest from .types.model_service import ListModelsResponse +from .types.model_service import ListModelVersionCheckpointsRequest +from .types.model_service import ListModelVersionCheckpointsResponse from .types.model_service import ListModelVersionsRequest from .types.model_service import ListModelVersionsResponse from .types.model_service import MergeVersionAliasesRequest +from .types.model_service import ModelVersionCheckpoint from .types.model_service import UpdateExplanationDatasetOperationMetadata from .types.model_service import UpdateExplanationDatasetRequest from .types.model_service import UpdateExplanationDatasetResponse @@ -755,6 +758,8 @@ from .types.notebook_service import UpgradeNotebookRuntimeRequest from .types.notebook_service import UpgradeNotebookRuntimeResponse from .types.notebook_service import NotebookExecutionJobView +from .types.notebook_software_config import NotebookSoftwareConfig +from .types.notebook_software_config import PostStartupScriptConfig from .types.openapi import Schema from .types.openapi import Type from .types.operation import DeleteOperationMetadata @@ -1534,6 +1539,8 @@ "ListModelEvaluationSlicesResponse", "ListModelEvaluationsRequest", "ListModelEvaluationsResponse", + "ListModelVersionCheckpointsRequest", + "ListModelVersionCheckpointsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", "ListModelsRequest", @@ -1620,6 +1627,7 @@ "ModelMonitoringStatsAnomalies", "ModelServiceClient", "ModelSourceInfo", + "ModelVersionCheckpoint", "MutateDeployedIndexOperationMetadata", "MutateDeployedIndexRequest", "MutateDeployedIndexResponse", @@ -1646,6 +1654,7 @@ "NotebookRuntimeTemplateRef", "NotebookRuntimeType", "NotebookServiceClient", + "NotebookSoftwareConfig", "PSCAutomationConfig", "PairwiseChoice", "PairwiseMetricInput", @@ -1679,6 +1688,7 @@ "PointwiseMetricResult", "PointwiseMetricSpec", "Port", + "PostStartupScriptConfig", "PredefinedSplit", "PredictRequest", "PredictRequestResponseLoggingConfig", diff --git a/google/cloud/aiplatform_v1/gapic_metadata.json b/google/cloud/aiplatform_v1/gapic_metadata.json index 1ccb108c52..170e56f293 100644 --- a/google/cloud/aiplatform_v1/gapic_metadata.json +++ b/google/cloud/aiplatform_v1/gapic_metadata.json @@ -3191,6 +3191,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" @@ -3286,6 +3291,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" @@ -3381,6 +3391,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" diff --git a/google/cloud/aiplatform_v1/gapic_version.py b/google/cloud/aiplatform_v1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform_v1/gapic_version.py +++ b/google/cloud/aiplatform_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform_v1/services/dataset_service/client.py b/google/cloud/aiplatform_v1/services/dataset_service/client.py index 36099764f8..d3de3520c6 100644 --- a/google/cloud/aiplatform_v1/services/dataset_service/client.py +++ b/google/cloud/aiplatform_v1/services/dataset_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -641,6 +643,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3248,16 +3277,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3303,16 +3336,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3475,16 +3512,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3596,16 +3637,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3718,16 +3763,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3778,16 +3827,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3833,16 +3886,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3888,16 +3945,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/dataset_service/transports/rest.py b/google/cloud/aiplatform_v1/services/dataset_service/transports/rest.py index 6cee292f17..dc17b99c5c 100644 --- a/google/cloud/aiplatform_v1/services/dataset_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/dataset_service/transports/rest.py @@ -261,12 +261,35 @@ def post_create_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset` interceptor runs + before the `post_create_dataset_with_metadata` interceptor. """ return response + def post_create_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_with_metadata` + interceptor in new development instead of the `post_create_dataset` interceptor. + When both interceptors are used, this `post_create_dataset_with_metadata` interceptor runs after the + `post_create_dataset` interceptor. The (possibly modified) response returned by + `post_create_dataset` will be passed to + `post_create_dataset_with_metadata`. + """ + return response, metadata + def pre_create_dataset_version( self, request: dataset_service.CreateDatasetVersionRequest, @@ -287,12 +310,35 @@ def post_create_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset_version` interceptor runs + before the `post_create_dataset_version_with_metadata` interceptor. """ return response + def post_create_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_version_with_metadata` + interceptor in new development instead of the `post_create_dataset_version` interceptor. + When both interceptors are used, this `post_create_dataset_version_with_metadata` interceptor runs after the + `post_create_dataset_version` interceptor. The (possibly modified) response returned by + `post_create_dataset_version` will be passed to + `post_create_dataset_version_with_metadata`. + """ + return response, metadata + def pre_delete_dataset( self, request: dataset_service.DeleteDatasetRequest, @@ -312,12 +358,35 @@ def post_delete_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset` interceptor runs + before the `post_delete_dataset_with_metadata` interceptor. """ return response + def post_delete_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_with_metadata` + interceptor in new development instead of the `post_delete_dataset` interceptor. + When both interceptors are used, this `post_delete_dataset_with_metadata` interceptor runs after the + `post_delete_dataset` interceptor. The (possibly modified) response returned by + `post_delete_dataset` will be passed to + `post_delete_dataset_with_metadata`. + """ + return response, metadata + def pre_delete_dataset_version( self, request: dataset_service.DeleteDatasetVersionRequest, @@ -338,12 +407,35 @@ def post_delete_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset_version` interceptor runs + before the `post_delete_dataset_version_with_metadata` interceptor. """ return response + def post_delete_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_version_with_metadata` + interceptor in new development instead of the `post_delete_dataset_version` interceptor. + When both interceptors are used, this `post_delete_dataset_version_with_metadata` interceptor runs after the + `post_delete_dataset_version` interceptor. The (possibly modified) response returned by + `post_delete_dataset_version` will be passed to + `post_delete_dataset_version_with_metadata`. + """ + return response, metadata + def pre_delete_saved_query( self, request: dataset_service.DeleteSavedQueryRequest, @@ -363,12 +455,35 @@ def post_delete_saved_query( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_saved_query - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_saved_query_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_saved_query` interceptor runs + before the `post_delete_saved_query_with_metadata` interceptor. """ return response + def post_delete_saved_query_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_saved_query + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_saved_query_with_metadata` + interceptor in new development instead of the `post_delete_saved_query` interceptor. + When both interceptors are used, this `post_delete_saved_query_with_metadata` interceptor runs after the + `post_delete_saved_query` interceptor. The (possibly modified) response returned by + `post_delete_saved_query` will be passed to + `post_delete_saved_query_with_metadata`. + """ + return response, metadata + def pre_export_data( self, request: dataset_service.ExportDataRequest, @@ -388,12 +503,35 @@ def post_export_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_export_data` interceptor runs + before the `post_export_data_with_metadata` interceptor. """ return response + def post_export_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_export_data_with_metadata` + interceptor in new development instead of the `post_export_data` interceptor. + When both interceptors are used, this `post_export_data_with_metadata` interceptor runs after the + `post_export_data` interceptor. The (possibly modified) response returned by + `post_export_data` will be passed to + `post_export_data_with_metadata`. + """ + return response, metadata + def pre_get_annotation_spec( self, request: dataset_service.GetAnnotationSpecRequest, @@ -414,12 +552,35 @@ def post_get_annotation_spec( ) -> annotation_spec.AnnotationSpec: """Post-rpc interceptor for get_annotation_spec - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_annotation_spec_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_annotation_spec` interceptor runs + before the `post_get_annotation_spec_with_metadata` interceptor. """ return response + def post_get_annotation_spec_with_metadata( + self, + response: annotation_spec.AnnotationSpec, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[annotation_spec.AnnotationSpec, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_annotation_spec + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_annotation_spec_with_metadata` + interceptor in new development instead of the `post_get_annotation_spec` interceptor. + When both interceptors are used, this `post_get_annotation_spec_with_metadata` interceptor runs after the + `post_get_annotation_spec` interceptor. The (possibly modified) response returned by + `post_get_annotation_spec` will be passed to + `post_get_annotation_spec_with_metadata`. + """ + return response, metadata + def pre_get_dataset( self, request: dataset_service.GetDatasetRequest, @@ -437,12 +598,35 @@ def pre_get_dataset( def post_get_dataset(self, response: dataset.Dataset) -> dataset.Dataset: """Post-rpc interceptor for get_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset` interceptor runs + before the `post_get_dataset_with_metadata` interceptor. """ return response + def post_get_dataset_with_metadata( + self, + response: dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_with_metadata` + interceptor in new development instead of the `post_get_dataset` interceptor. + When both interceptors are used, this `post_get_dataset_with_metadata` interceptor runs after the + `post_get_dataset` interceptor. The (possibly modified) response returned by + `post_get_dataset` will be passed to + `post_get_dataset_with_metadata`. + """ + return response, metadata + def pre_get_dataset_version( self, request: dataset_service.GetDatasetVersionRequest, @@ -463,12 +647,35 @@ def post_get_dataset_version( ) -> dataset_version.DatasetVersion: """Post-rpc interceptor for get_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset_version` interceptor runs + before the `post_get_dataset_version_with_metadata` interceptor. """ return response + def post_get_dataset_version_with_metadata( + self, + response: dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_version_with_metadata` + interceptor in new development instead of the `post_get_dataset_version` interceptor. + When both interceptors are used, this `post_get_dataset_version_with_metadata` interceptor runs after the + `post_get_dataset_version` interceptor. The (possibly modified) response returned by + `post_get_dataset_version` will be passed to + `post_get_dataset_version_with_metadata`. + """ + return response, metadata + def pre_import_data( self, request: dataset_service.ImportDataRequest, @@ -488,12 +695,35 @@ def post_import_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_import_data` interceptor runs + before the `post_import_data_with_metadata` interceptor. """ return response + def post_import_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_import_data_with_metadata` + interceptor in new development instead of the `post_import_data` interceptor. + When both interceptors are used, this `post_import_data_with_metadata` interceptor runs after the + `post_import_data` interceptor. The (possibly modified) response returned by + `post_import_data` will be passed to + `post_import_data_with_metadata`. + """ + return response, metadata + def pre_list_annotations( self, request: dataset_service.ListAnnotationsRequest, @@ -513,12 +743,37 @@ def post_list_annotations( ) -> dataset_service.ListAnnotationsResponse: """Post-rpc interceptor for list_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_annotations` interceptor runs + before the `post_list_annotations_with_metadata` interceptor. """ return response + def post_list_annotations_with_metadata( + self, + response: dataset_service.ListAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListAnnotationsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_annotations_with_metadata` + interceptor in new development instead of the `post_list_annotations` interceptor. + When both interceptors are used, this `post_list_annotations_with_metadata` interceptor runs after the + `post_list_annotations` interceptor. The (possibly modified) response returned by + `post_list_annotations` will be passed to + `post_list_annotations_with_metadata`. + """ + return response, metadata + def pre_list_data_items( self, request: dataset_service.ListDataItemsRequest, @@ -538,12 +793,37 @@ def post_list_data_items( ) -> dataset_service.ListDataItemsResponse: """Post-rpc interceptor for list_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_items` interceptor runs + before the `post_list_data_items_with_metadata` interceptor. """ return response + def post_list_data_items_with_metadata( + self, + response: dataset_service.ListDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_data_items_with_metadata` + interceptor in new development instead of the `post_list_data_items` interceptor. + When both interceptors are used, this `post_list_data_items_with_metadata` interceptor runs after the + `post_list_data_items` interceptor. The (possibly modified) response returned by + `post_list_data_items` will be passed to + `post_list_data_items_with_metadata`. + """ + return response, metadata + def pre_list_datasets( self, request: dataset_service.ListDatasetsRequest, @@ -563,12 +843,37 @@ def post_list_datasets( ) -> dataset_service.ListDatasetsResponse: """Post-rpc interceptor for list_datasets - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_datasets_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_datasets` interceptor runs + before the `post_list_datasets_with_metadata` interceptor. """ return response + def post_list_datasets_with_metadata( + self, + response: dataset_service.ListDatasetsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_datasets + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_datasets_with_metadata` + interceptor in new development instead of the `post_list_datasets` interceptor. + When both interceptors are used, this `post_list_datasets_with_metadata` interceptor runs after the + `post_list_datasets` interceptor. The (possibly modified) response returned by + `post_list_datasets` will be passed to + `post_list_datasets_with_metadata`. + """ + return response, metadata + def pre_list_dataset_versions( self, request: dataset_service.ListDatasetVersionsRequest, @@ -589,12 +894,38 @@ def post_list_dataset_versions( ) -> dataset_service.ListDatasetVersionsResponse: """Post-rpc interceptor for list_dataset_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_dataset_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_dataset_versions` interceptor runs + before the `post_list_dataset_versions_with_metadata` interceptor. """ return response + def post_list_dataset_versions_with_metadata( + self, + response: dataset_service.ListDatasetVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetVersionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_dataset_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_dataset_versions_with_metadata` + interceptor in new development instead of the `post_list_dataset_versions` interceptor. + When both interceptors are used, this `post_list_dataset_versions_with_metadata` interceptor runs after the + `post_list_dataset_versions` interceptor. The (possibly modified) response returned by + `post_list_dataset_versions` will be passed to + `post_list_dataset_versions_with_metadata`. + """ + return response, metadata + def pre_list_saved_queries( self, request: dataset_service.ListSavedQueriesRequest, @@ -614,12 +945,38 @@ def post_list_saved_queries( ) -> dataset_service.ListSavedQueriesResponse: """Post-rpc interceptor for list_saved_queries - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_saved_queries_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_saved_queries` interceptor runs + before the `post_list_saved_queries_with_metadata` interceptor. """ return response + def post_list_saved_queries_with_metadata( + self, + response: dataset_service.ListSavedQueriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListSavedQueriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_saved_queries + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_saved_queries_with_metadata` + interceptor in new development instead of the `post_list_saved_queries` interceptor. + When both interceptors are used, this `post_list_saved_queries_with_metadata` interceptor runs after the + `post_list_saved_queries` interceptor. The (possibly modified) response returned by + `post_list_saved_queries` will be passed to + `post_list_saved_queries_with_metadata`. + """ + return response, metadata + def pre_restore_dataset_version( self, request: dataset_service.RestoreDatasetVersionRequest, @@ -640,12 +997,35 @@ def post_restore_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for restore_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_restore_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_restore_dataset_version` interceptor runs + before the `post_restore_dataset_version_with_metadata` interceptor. """ return response + def post_restore_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for restore_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_restore_dataset_version_with_metadata` + interceptor in new development instead of the `post_restore_dataset_version` interceptor. + When both interceptors are used, this `post_restore_dataset_version_with_metadata` interceptor runs after the + `post_restore_dataset_version` interceptor. The (possibly modified) response returned by + `post_restore_dataset_version` will be passed to + `post_restore_dataset_version_with_metadata`. + """ + return response, metadata + def pre_search_data_items( self, request: dataset_service.SearchDataItemsRequest, @@ -665,12 +1045,37 @@ def post_search_data_items( ) -> dataset_service.SearchDataItemsResponse: """Post-rpc interceptor for search_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_search_data_items` interceptor runs + before the `post_search_data_items_with_metadata` interceptor. """ return response + def post_search_data_items_with_metadata( + self, + response: dataset_service.SearchDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.SearchDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for search_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_search_data_items_with_metadata` + interceptor in new development instead of the `post_search_data_items` interceptor. + When both interceptors are used, this `post_search_data_items_with_metadata` interceptor runs after the + `post_search_data_items` interceptor. The (possibly modified) response returned by + `post_search_data_items` will be passed to + `post_search_data_items_with_metadata`. + """ + return response, metadata + def pre_update_dataset( self, request: dataset_service.UpdateDatasetRequest, @@ -688,12 +1093,35 @@ def pre_update_dataset( def post_update_dataset(self, response: gca_dataset.Dataset) -> gca_dataset.Dataset: """Post-rpc interceptor for update_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset` interceptor runs + before the `post_update_dataset_with_metadata` interceptor. """ return response + def post_update_dataset_with_metadata( + self, + response: gca_dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_with_metadata` + interceptor in new development instead of the `post_update_dataset` interceptor. + When both interceptors are used, this `post_update_dataset_with_metadata` interceptor runs after the + `post_update_dataset` interceptor. The (possibly modified) response returned by + `post_update_dataset` will be passed to + `post_update_dataset_with_metadata`. + """ + return response, metadata + def pre_update_dataset_version( self, request: dataset_service.UpdateDatasetVersionRequest, @@ -714,12 +1142,37 @@ def post_update_dataset_version( ) -> gca_dataset_version.DatasetVersion: """Post-rpc interceptor for update_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset_version` interceptor runs + before the `post_update_dataset_version_with_metadata` interceptor. """ return response + def post_update_dataset_version_with_metadata( + self, + response: gca_dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_version_with_metadata` + interceptor in new development instead of the `post_update_dataset_version` interceptor. + When both interceptors are used, this `post_update_dataset_version_with_metadata` interceptor runs after the + `post_update_dataset_version` interceptor. The (possibly modified) response returned by + `post_update_dataset_version` will be passed to + `post_update_dataset_version_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3092,6 +3545,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3244,6 +3701,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3387,6 +3848,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3533,6 +3998,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3678,6 +4147,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_saved_query(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_saved_query_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3831,6 +4304,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3977,6 +4454,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_annotation_spec(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_annotation_spec_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4124,6 +4605,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4269,6 +4754,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4422,6 +4911,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4568,6 +5061,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4714,6 +5211,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4860,6 +5361,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_datasets(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_datasets_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5009,6 +5514,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_dataset_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_dataset_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5157,6 +5666,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_saved_queries(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_saved_queries_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5305,6 +5818,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_restore_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_restore_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5451,6 +5968,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5603,6 +6124,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5754,6 +6279,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/dataset_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/dataset_service/transports/rest_asyncio.py index 3426f06746..2375404c68 100644 --- a/google/cloud/aiplatform_v1/services/dataset_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/dataset_service/transports/rest_asyncio.py @@ -279,12 +279,35 @@ async def post_create_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset` interceptor runs + before the `post_create_dataset_with_metadata` interceptor. """ return response + async def post_create_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_with_metadata` + interceptor in new development instead of the `post_create_dataset` interceptor. + When both interceptors are used, this `post_create_dataset_with_metadata` interceptor runs after the + `post_create_dataset` interceptor. The (possibly modified) response returned by + `post_create_dataset` will be passed to + `post_create_dataset_with_metadata`. + """ + return response, metadata + async def pre_create_dataset_version( self, request: dataset_service.CreateDatasetVersionRequest, @@ -305,12 +328,35 @@ async def post_create_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset_version` interceptor runs + before the `post_create_dataset_version_with_metadata` interceptor. """ return response + async def post_create_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_version_with_metadata` + interceptor in new development instead of the `post_create_dataset_version` interceptor. + When both interceptors are used, this `post_create_dataset_version_with_metadata` interceptor runs after the + `post_create_dataset_version` interceptor. The (possibly modified) response returned by + `post_create_dataset_version` will be passed to + `post_create_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_delete_dataset( self, request: dataset_service.DeleteDatasetRequest, @@ -330,12 +376,35 @@ async def post_delete_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset` interceptor runs + before the `post_delete_dataset_with_metadata` interceptor. """ return response + async def post_delete_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_with_metadata` + interceptor in new development instead of the `post_delete_dataset` interceptor. + When both interceptors are used, this `post_delete_dataset_with_metadata` interceptor runs after the + `post_delete_dataset` interceptor. The (possibly modified) response returned by + `post_delete_dataset` will be passed to + `post_delete_dataset_with_metadata`. + """ + return response, metadata + async def pre_delete_dataset_version( self, request: dataset_service.DeleteDatasetVersionRequest, @@ -356,12 +425,35 @@ async def post_delete_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset_version` interceptor runs + before the `post_delete_dataset_version_with_metadata` interceptor. """ return response + async def post_delete_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_version_with_metadata` + interceptor in new development instead of the `post_delete_dataset_version` interceptor. + When both interceptors are used, this `post_delete_dataset_version_with_metadata` interceptor runs after the + `post_delete_dataset_version` interceptor. The (possibly modified) response returned by + `post_delete_dataset_version` will be passed to + `post_delete_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_delete_saved_query( self, request: dataset_service.DeleteSavedQueryRequest, @@ -381,12 +473,35 @@ async def post_delete_saved_query( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_saved_query - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_saved_query_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_saved_query` interceptor runs + before the `post_delete_saved_query_with_metadata` interceptor. """ return response + async def post_delete_saved_query_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_saved_query + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_saved_query_with_metadata` + interceptor in new development instead of the `post_delete_saved_query` interceptor. + When both interceptors are used, this `post_delete_saved_query_with_metadata` interceptor runs after the + `post_delete_saved_query` interceptor. The (possibly modified) response returned by + `post_delete_saved_query` will be passed to + `post_delete_saved_query_with_metadata`. + """ + return response, metadata + async def pre_export_data( self, request: dataset_service.ExportDataRequest, @@ -406,12 +521,35 @@ async def post_export_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_export_data` interceptor runs + before the `post_export_data_with_metadata` interceptor. """ return response + async def post_export_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_export_data_with_metadata` + interceptor in new development instead of the `post_export_data` interceptor. + When both interceptors are used, this `post_export_data_with_metadata` interceptor runs after the + `post_export_data` interceptor. The (possibly modified) response returned by + `post_export_data` will be passed to + `post_export_data_with_metadata`. + """ + return response, metadata + async def pre_get_annotation_spec( self, request: dataset_service.GetAnnotationSpecRequest, @@ -432,12 +570,35 @@ async def post_get_annotation_spec( ) -> annotation_spec.AnnotationSpec: """Post-rpc interceptor for get_annotation_spec - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_annotation_spec_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_annotation_spec` interceptor runs + before the `post_get_annotation_spec_with_metadata` interceptor. """ return response + async def post_get_annotation_spec_with_metadata( + self, + response: annotation_spec.AnnotationSpec, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[annotation_spec.AnnotationSpec, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_annotation_spec + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_annotation_spec_with_metadata` + interceptor in new development instead of the `post_get_annotation_spec` interceptor. + When both interceptors are used, this `post_get_annotation_spec_with_metadata` interceptor runs after the + `post_get_annotation_spec` interceptor. The (possibly modified) response returned by + `post_get_annotation_spec` will be passed to + `post_get_annotation_spec_with_metadata`. + """ + return response, metadata + async def pre_get_dataset( self, request: dataset_service.GetDatasetRequest, @@ -455,12 +616,35 @@ async def pre_get_dataset( async def post_get_dataset(self, response: dataset.Dataset) -> dataset.Dataset: """Post-rpc interceptor for get_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset` interceptor runs + before the `post_get_dataset_with_metadata` interceptor. """ return response + async def post_get_dataset_with_metadata( + self, + response: dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_with_metadata` + interceptor in new development instead of the `post_get_dataset` interceptor. + When both interceptors are used, this `post_get_dataset_with_metadata` interceptor runs after the + `post_get_dataset` interceptor. The (possibly modified) response returned by + `post_get_dataset` will be passed to + `post_get_dataset_with_metadata`. + """ + return response, metadata + async def pre_get_dataset_version( self, request: dataset_service.GetDatasetVersionRequest, @@ -481,12 +665,35 @@ async def post_get_dataset_version( ) -> dataset_version.DatasetVersion: """Post-rpc interceptor for get_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset_version` interceptor runs + before the `post_get_dataset_version_with_metadata` interceptor. """ return response + async def post_get_dataset_version_with_metadata( + self, + response: dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_version_with_metadata` + interceptor in new development instead of the `post_get_dataset_version` interceptor. + When both interceptors are used, this `post_get_dataset_version_with_metadata` interceptor runs after the + `post_get_dataset_version` interceptor. The (possibly modified) response returned by + `post_get_dataset_version` will be passed to + `post_get_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_import_data( self, request: dataset_service.ImportDataRequest, @@ -506,12 +713,35 @@ async def post_import_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_import_data` interceptor runs + before the `post_import_data_with_metadata` interceptor. """ return response + async def post_import_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_import_data_with_metadata` + interceptor in new development instead of the `post_import_data` interceptor. + When both interceptors are used, this `post_import_data_with_metadata` interceptor runs after the + `post_import_data` interceptor. The (possibly modified) response returned by + `post_import_data` will be passed to + `post_import_data_with_metadata`. + """ + return response, metadata + async def pre_list_annotations( self, request: dataset_service.ListAnnotationsRequest, @@ -531,12 +761,37 @@ async def post_list_annotations( ) -> dataset_service.ListAnnotationsResponse: """Post-rpc interceptor for list_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_annotations` interceptor runs + before the `post_list_annotations_with_metadata` interceptor. """ return response + async def post_list_annotations_with_metadata( + self, + response: dataset_service.ListAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListAnnotationsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_annotations_with_metadata` + interceptor in new development instead of the `post_list_annotations` interceptor. + When both interceptors are used, this `post_list_annotations_with_metadata` interceptor runs after the + `post_list_annotations` interceptor. The (possibly modified) response returned by + `post_list_annotations` will be passed to + `post_list_annotations_with_metadata`. + """ + return response, metadata + async def pre_list_data_items( self, request: dataset_service.ListDataItemsRequest, @@ -556,12 +811,37 @@ async def post_list_data_items( ) -> dataset_service.ListDataItemsResponse: """Post-rpc interceptor for list_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_items` interceptor runs + before the `post_list_data_items_with_metadata` interceptor. """ return response + async def post_list_data_items_with_metadata( + self, + response: dataset_service.ListDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_data_items_with_metadata` + interceptor in new development instead of the `post_list_data_items` interceptor. + When both interceptors are used, this `post_list_data_items_with_metadata` interceptor runs after the + `post_list_data_items` interceptor. The (possibly modified) response returned by + `post_list_data_items` will be passed to + `post_list_data_items_with_metadata`. + """ + return response, metadata + async def pre_list_datasets( self, request: dataset_service.ListDatasetsRequest, @@ -581,12 +861,37 @@ async def post_list_datasets( ) -> dataset_service.ListDatasetsResponse: """Post-rpc interceptor for list_datasets - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_datasets_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_datasets` interceptor runs + before the `post_list_datasets_with_metadata` interceptor. """ return response + async def post_list_datasets_with_metadata( + self, + response: dataset_service.ListDatasetsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_datasets + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_datasets_with_metadata` + interceptor in new development instead of the `post_list_datasets` interceptor. + When both interceptors are used, this `post_list_datasets_with_metadata` interceptor runs after the + `post_list_datasets` interceptor. The (possibly modified) response returned by + `post_list_datasets` will be passed to + `post_list_datasets_with_metadata`. + """ + return response, metadata + async def pre_list_dataset_versions( self, request: dataset_service.ListDatasetVersionsRequest, @@ -607,12 +912,38 @@ async def post_list_dataset_versions( ) -> dataset_service.ListDatasetVersionsResponse: """Post-rpc interceptor for list_dataset_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_dataset_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_dataset_versions` interceptor runs + before the `post_list_dataset_versions_with_metadata` interceptor. """ return response + async def post_list_dataset_versions_with_metadata( + self, + response: dataset_service.ListDatasetVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetVersionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_dataset_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_dataset_versions_with_metadata` + interceptor in new development instead of the `post_list_dataset_versions` interceptor. + When both interceptors are used, this `post_list_dataset_versions_with_metadata` interceptor runs after the + `post_list_dataset_versions` interceptor. The (possibly modified) response returned by + `post_list_dataset_versions` will be passed to + `post_list_dataset_versions_with_metadata`. + """ + return response, metadata + async def pre_list_saved_queries( self, request: dataset_service.ListSavedQueriesRequest, @@ -632,12 +963,38 @@ async def post_list_saved_queries( ) -> dataset_service.ListSavedQueriesResponse: """Post-rpc interceptor for list_saved_queries - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_saved_queries_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_saved_queries` interceptor runs + before the `post_list_saved_queries_with_metadata` interceptor. """ return response + async def post_list_saved_queries_with_metadata( + self, + response: dataset_service.ListSavedQueriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListSavedQueriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_saved_queries + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_saved_queries_with_metadata` + interceptor in new development instead of the `post_list_saved_queries` interceptor. + When both interceptors are used, this `post_list_saved_queries_with_metadata` interceptor runs after the + `post_list_saved_queries` interceptor. The (possibly modified) response returned by + `post_list_saved_queries` will be passed to + `post_list_saved_queries_with_metadata`. + """ + return response, metadata + async def pre_restore_dataset_version( self, request: dataset_service.RestoreDatasetVersionRequest, @@ -658,12 +1015,35 @@ async def post_restore_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for restore_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_restore_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_restore_dataset_version` interceptor runs + before the `post_restore_dataset_version_with_metadata` interceptor. """ return response + async def post_restore_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for restore_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_restore_dataset_version_with_metadata` + interceptor in new development instead of the `post_restore_dataset_version` interceptor. + When both interceptors are used, this `post_restore_dataset_version_with_metadata` interceptor runs after the + `post_restore_dataset_version` interceptor. The (possibly modified) response returned by + `post_restore_dataset_version` will be passed to + `post_restore_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_search_data_items( self, request: dataset_service.SearchDataItemsRequest, @@ -683,12 +1063,37 @@ async def post_search_data_items( ) -> dataset_service.SearchDataItemsResponse: """Post-rpc interceptor for search_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_search_data_items` interceptor runs + before the `post_search_data_items_with_metadata` interceptor. """ return response + async def post_search_data_items_with_metadata( + self, + response: dataset_service.SearchDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.SearchDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for search_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_search_data_items_with_metadata` + interceptor in new development instead of the `post_search_data_items` interceptor. + When both interceptors are used, this `post_search_data_items_with_metadata` interceptor runs after the + `post_search_data_items` interceptor. The (possibly modified) response returned by + `post_search_data_items` will be passed to + `post_search_data_items_with_metadata`. + """ + return response, metadata + async def pre_update_dataset( self, request: dataset_service.UpdateDatasetRequest, @@ -708,12 +1113,35 @@ async def post_update_dataset( ) -> gca_dataset.Dataset: """Post-rpc interceptor for update_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset` interceptor runs + before the `post_update_dataset_with_metadata` interceptor. """ return response + async def post_update_dataset_with_metadata( + self, + response: gca_dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_with_metadata` + interceptor in new development instead of the `post_update_dataset` interceptor. + When both interceptors are used, this `post_update_dataset_with_metadata` interceptor runs after the + `post_update_dataset` interceptor. The (possibly modified) response returned by + `post_update_dataset` will be passed to + `post_update_dataset_with_metadata`. + """ + return response, metadata + async def pre_update_dataset_version( self, request: dataset_service.UpdateDatasetVersionRequest, @@ -734,12 +1162,37 @@ async def post_update_dataset_version( ) -> gca_dataset_version.DatasetVersion: """Post-rpc interceptor for update_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset_version` interceptor runs + before the `post_update_dataset_version_with_metadata` interceptor. """ return response + async def post_update_dataset_version_with_metadata( + self, + response: gca_dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_version_with_metadata` + interceptor in new development instead of the `post_update_dataset_version` interceptor. + When both interceptors are used, this `post_update_dataset_version_with_metadata` interceptor runs after the + `post_update_dataset_version` interceptor. The (possibly modified) response returned by + `post_update_dataset_version` will be passed to + `post_update_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1345,6 +1798,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1505,6 +1962,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1660,6 +2121,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1814,6 +2279,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1970,6 +2439,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_saved_query(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_saved_query_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2133,6 +2606,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2288,6 +2765,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_annotation_spec(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_annotation_spec_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2443,6 +2924,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2597,6 +3082,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2760,6 +3249,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2915,6 +3408,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3071,6 +3568,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3227,6 +3728,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_datasets(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_datasets_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3382,6 +3887,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_dataset_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_dataset_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3539,6 +4048,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_saved_queries(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_saved_queries_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3695,6 +4208,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_restore_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_restore_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3850,6 +4370,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_search_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4012,6 +4536,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4169,6 +4697,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/client.py b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/client.py index 13313ad94e..00f680e82e 100644 --- a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/client.py +++ b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -588,6 +590,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1730,16 +1759,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1785,16 +1818,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1957,16 +1994,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2078,16 +2119,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2200,16 +2245,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2260,16 +2309,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2315,16 +2368,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2370,16 +2427,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest.py b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest.py index 87ce8a74d1..1e50e2f380 100644 --- a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest.py @@ -154,12 +154,35 @@ def post_create_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_deployment_resource_pool` interceptor runs + before the `post_create_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_create_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_create_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_create_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_create_deployment_resource_pool_with_metadata` interceptor runs after the + `post_create_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_create_deployment_resource_pool` will be passed to + `post_create_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_delete_deployment_resource_pool( self, request: deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest, @@ -180,12 +203,35 @@ def post_delete_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_deployment_resource_pool` interceptor runs + before the `post_delete_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_delete_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_delete_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_delete_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_delete_deployment_resource_pool_with_metadata` interceptor runs after the + `post_delete_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_delete_deployment_resource_pool` will be passed to + `post_delete_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_get_deployment_resource_pool( self, request: deployment_resource_pool_service.GetDeploymentResourcePoolRequest, @@ -206,12 +252,38 @@ def post_get_deployment_resource_pool( ) -> deployment_resource_pool.DeploymentResourcePool: """Post-rpc interceptor for get_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_deployment_resource_pool` interceptor runs + before the `post_get_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_get_deployment_resource_pool_with_metadata( + self, + response: deployment_resource_pool.DeploymentResourcePool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool.DeploymentResourcePool, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_get_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_get_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_get_deployment_resource_pool_with_metadata` interceptor runs after the + `post_get_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_get_deployment_resource_pool` will be passed to + `post_get_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_list_deployment_resource_pools( self, request: deployment_resource_pool_service.ListDeploymentResourcePoolsRequest, @@ -233,12 +305,38 @@ def post_list_deployment_resource_pools( ) -> deployment_resource_pool_service.ListDeploymentResourcePoolsResponse: """Post-rpc interceptor for list_deployment_resource_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_deployment_resource_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_deployment_resource_pools` interceptor runs + before the `post_list_deployment_resource_pools_with_metadata` interceptor. """ return response + def post_list_deployment_resource_pools_with_metadata( + self, + response: deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_deployment_resource_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_list_deployment_resource_pools_with_metadata` + interceptor in new development instead of the `post_list_deployment_resource_pools` interceptor. + When both interceptors are used, this `post_list_deployment_resource_pools_with_metadata` interceptor runs after the + `post_list_deployment_resource_pools` interceptor. The (possibly modified) response returned by + `post_list_deployment_resource_pools` will be passed to + `post_list_deployment_resource_pools_with_metadata`. + """ + return response, metadata + def pre_query_deployed_models( self, request: deployment_resource_pool_service.QueryDeployedModelsRequest, @@ -259,12 +357,38 @@ def post_query_deployed_models( ) -> deployment_resource_pool_service.QueryDeployedModelsResponse: """Post-rpc interceptor for query_deployed_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_deployed_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_query_deployed_models` interceptor runs + before the `post_query_deployed_models_with_metadata` interceptor. """ return response + def post_query_deployed_models_with_metadata( + self, + response: deployment_resource_pool_service.QueryDeployedModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.QueryDeployedModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_deployed_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_query_deployed_models_with_metadata` + interceptor in new development instead of the `post_query_deployed_models` interceptor. + When both interceptors are used, this `post_query_deployed_models_with_metadata` interceptor runs after the + `post_query_deployed_models` interceptor. The (possibly modified) response returned by + `post_query_deployed_models` will be passed to + `post_query_deployed_models_with_metadata`. + """ + return response, metadata + def pre_update_deployment_resource_pool( self, request: deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest, @@ -285,12 +409,35 @@ def post_update_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_deployment_resource_pool` interceptor runs + before the `post_update_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_update_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_update_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_update_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_update_deployment_resource_pool_with_metadata` interceptor runs after the + `post_update_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_update_deployment_resource_pool` will be passed to + `post_update_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2672,6 +2819,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2821,6 +2975,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2973,6 +3134,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3133,6 +3298,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_deployment_resource_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_deployment_resource_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3286,6 +3458,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_deployed_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_query_deployed_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3443,6 +3619,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest_asyncio.py index 9cf7e5cf51..d799b81929 100644 --- a/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/deployment_resource_pool_service/transports/rest_asyncio.py @@ -172,12 +172,35 @@ async def post_create_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_deployment_resource_pool` interceptor runs + before the `post_create_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_create_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_create_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_create_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_create_deployment_resource_pool_with_metadata` interceptor runs after the + `post_create_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_create_deployment_resource_pool` will be passed to + `post_create_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_delete_deployment_resource_pool( self, request: deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest, @@ -198,12 +221,35 @@ async def post_delete_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_deployment_resource_pool` interceptor runs + before the `post_delete_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_delete_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_delete_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_delete_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_delete_deployment_resource_pool_with_metadata` interceptor runs after the + `post_delete_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_delete_deployment_resource_pool` will be passed to + `post_delete_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_get_deployment_resource_pool( self, request: deployment_resource_pool_service.GetDeploymentResourcePoolRequest, @@ -224,12 +270,38 @@ async def post_get_deployment_resource_pool( ) -> deployment_resource_pool.DeploymentResourcePool: """Post-rpc interceptor for get_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_deployment_resource_pool` interceptor runs + before the `post_get_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_get_deployment_resource_pool_with_metadata( + self, + response: deployment_resource_pool.DeploymentResourcePool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool.DeploymentResourcePool, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_get_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_get_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_get_deployment_resource_pool_with_metadata` interceptor runs after the + `post_get_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_get_deployment_resource_pool` will be passed to + `post_get_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_list_deployment_resource_pools( self, request: deployment_resource_pool_service.ListDeploymentResourcePoolsRequest, @@ -251,12 +323,38 @@ async def post_list_deployment_resource_pools( ) -> deployment_resource_pool_service.ListDeploymentResourcePoolsResponse: """Post-rpc interceptor for list_deployment_resource_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_deployment_resource_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_deployment_resource_pools` interceptor runs + before the `post_list_deployment_resource_pools_with_metadata` interceptor. """ return response + async def post_list_deployment_resource_pools_with_metadata( + self, + response: deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_deployment_resource_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_list_deployment_resource_pools_with_metadata` + interceptor in new development instead of the `post_list_deployment_resource_pools` interceptor. + When both interceptors are used, this `post_list_deployment_resource_pools_with_metadata` interceptor runs after the + `post_list_deployment_resource_pools` interceptor. The (possibly modified) response returned by + `post_list_deployment_resource_pools` will be passed to + `post_list_deployment_resource_pools_with_metadata`. + """ + return response, metadata + async def pre_query_deployed_models( self, request: deployment_resource_pool_service.QueryDeployedModelsRequest, @@ -277,12 +375,38 @@ async def post_query_deployed_models( ) -> deployment_resource_pool_service.QueryDeployedModelsResponse: """Post-rpc interceptor for query_deployed_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_deployed_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_query_deployed_models` interceptor runs + before the `post_query_deployed_models_with_metadata` interceptor. """ return response + async def post_query_deployed_models_with_metadata( + self, + response: deployment_resource_pool_service.QueryDeployedModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.QueryDeployedModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_deployed_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_query_deployed_models_with_metadata` + interceptor in new development instead of the `post_query_deployed_models` interceptor. + When both interceptors are used, this `post_query_deployed_models_with_metadata` interceptor runs after the + `post_query_deployed_models` interceptor. The (possibly modified) response returned by + `post_query_deployed_models` will be passed to + `post_query_deployed_models_with_metadata`. + """ + return response, metadata + async def pre_update_deployment_resource_pool( self, request: deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest, @@ -303,12 +427,35 @@ async def post_update_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_deployment_resource_pool` interceptor runs + before the `post_update_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_update_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_update_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_update_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_update_deployment_resource_pool_with_metadata` interceptor runs after the + `post_update_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_update_deployment_resource_pool` will be passed to + `post_update_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -857,6 +1004,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1017,6 +1171,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1178,6 +1339,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1347,6 +1515,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_deployment_resource_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_deployment_resource_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1506,6 +1681,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_deployed_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_query_deployed_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1674,6 +1853,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/endpoint_service/client.py b/google/cloud/aiplatform_v1/services/endpoint_service/client.py index 8185697f63..f95d54248c 100644 --- a/google/cloud/aiplatform_v1/services/endpoint_service/client.py +++ b/google/cloud/aiplatform_v1/services/endpoint_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -618,6 +620,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2137,16 +2166,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2192,16 +2225,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2364,16 +2401,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2485,16 +2526,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2607,16 +2652,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2667,16 +2716,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2722,16 +2775,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2777,16 +2834,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest.py b/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest.py index 8f0c05c205..db8126d2a6 100644 --- a/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest.py @@ -178,12 +178,35 @@ def post_create_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_endpoint` interceptor runs + before the `post_create_endpoint_with_metadata` interceptor. """ return response + def post_create_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_create_endpoint_with_metadata` + interceptor in new development instead of the `post_create_endpoint` interceptor. + When both interceptors are used, this `post_create_endpoint_with_metadata` interceptor runs after the + `post_create_endpoint` interceptor. The (possibly modified) response returned by + `post_create_endpoint` will be passed to + `post_create_endpoint_with_metadata`. + """ + return response, metadata + def pre_delete_endpoint( self, request: endpoint_service.DeleteEndpointRequest, @@ -203,12 +226,35 @@ def post_delete_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_endpoint` interceptor runs + before the `post_delete_endpoint_with_metadata` interceptor. """ return response + def post_delete_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_endpoint` interceptor. + When both interceptors are used, this `post_delete_endpoint_with_metadata` interceptor runs after the + `post_delete_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_endpoint` will be passed to + `post_delete_endpoint_with_metadata`. + """ + return response, metadata + def pre_deploy_model( self, request: endpoint_service.DeployModelRequest, @@ -228,12 +274,35 @@ def post_deploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_model` interceptor runs + before the `post_deploy_model_with_metadata` interceptor. """ return response + def post_deploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_model_with_metadata` + interceptor in new development instead of the `post_deploy_model` interceptor. + When both interceptors are used, this `post_deploy_model_with_metadata` interceptor runs after the + `post_deploy_model` interceptor. The (possibly modified) response returned by + `post_deploy_model` will be passed to + `post_deploy_model_with_metadata`. + """ + return response, metadata + def pre_get_endpoint( self, request: endpoint_service.GetEndpointRequest, @@ -251,12 +320,35 @@ def pre_get_endpoint( def post_get_endpoint(self, response: endpoint.Endpoint) -> endpoint.Endpoint: """Post-rpc interceptor for get_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_endpoint` interceptor runs + before the `post_get_endpoint_with_metadata` interceptor. """ return response + def post_get_endpoint_with_metadata( + self, + response: endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_get_endpoint_with_metadata` + interceptor in new development instead of the `post_get_endpoint` interceptor. + When both interceptors are used, this `post_get_endpoint_with_metadata` interceptor runs after the + `post_get_endpoint` interceptor. The (possibly modified) response returned by + `post_get_endpoint` will be passed to + `post_get_endpoint_with_metadata`. + """ + return response, metadata + def pre_list_endpoints( self, request: endpoint_service.ListEndpointsRequest, @@ -276,12 +368,37 @@ def post_list_endpoints( ) -> endpoint_service.ListEndpointsResponse: """Post-rpc interceptor for list_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_endpoints` interceptor runs + before the `post_list_endpoints_with_metadata` interceptor. """ return response + def post_list_endpoints_with_metadata( + self, + response: endpoint_service.ListEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + endpoint_service.ListEndpointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_list_endpoints_with_metadata` + interceptor in new development instead of the `post_list_endpoints` interceptor. + When both interceptors are used, this `post_list_endpoints_with_metadata` interceptor runs after the + `post_list_endpoints` interceptor. The (possibly modified) response returned by + `post_list_endpoints` will be passed to + `post_list_endpoints_with_metadata`. + """ + return response, metadata + def pre_mutate_deployed_model( self, request: endpoint_service.MutateDeployedModelRequest, @@ -302,12 +419,35 @@ def post_mutate_deployed_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_model` interceptor runs + before the `post_mutate_deployed_model_with_metadata` interceptor. """ return response + def post_mutate_deployed_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_model_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_model` interceptor. + When both interceptors are used, this `post_mutate_deployed_model_with_metadata` interceptor runs after the + `post_mutate_deployed_model` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_model` will be passed to + `post_mutate_deployed_model_with_metadata`. + """ + return response, metadata + def pre_undeploy_model( self, request: endpoint_service.UndeployModelRequest, @@ -327,12 +467,35 @@ def post_undeploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_model` interceptor runs + before the `post_undeploy_model_with_metadata` interceptor. """ return response + def post_undeploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_model_with_metadata` + interceptor in new development instead of the `post_undeploy_model` interceptor. + When both interceptors are used, this `post_undeploy_model_with_metadata` interceptor runs after the + `post_undeploy_model` interceptor. The (possibly modified) response returned by + `post_undeploy_model` will be passed to + `post_undeploy_model_with_metadata`. + """ + return response, metadata + def pre_update_endpoint( self, request: endpoint_service.UpdateEndpointRequest, @@ -352,12 +515,35 @@ def post_update_endpoint( ) -> gca_endpoint.Endpoint: """Post-rpc interceptor for update_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint` interceptor runs + before the `post_update_endpoint_with_metadata` interceptor. """ return response + def post_update_endpoint_with_metadata( + self, + response: gca_endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_with_metadata` + interceptor in new development instead of the `post_update_endpoint` interceptor. + When both interceptors are used, this `post_update_endpoint_with_metadata` interceptor runs after the + `post_update_endpoint` interceptor. The (possibly modified) response returned by + `post_update_endpoint` will be passed to + `post_update_endpoint_with_metadata`. + """ + return response, metadata + def pre_update_endpoint_long_running( self, request: endpoint_service.UpdateEndpointLongRunningRequest, @@ -378,12 +564,35 @@ def post_update_endpoint_long_running( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_endpoint_long_running - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_long_running_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint_long_running` interceptor runs + before the `post_update_endpoint_long_running_with_metadata` interceptor. """ return response + def post_update_endpoint_long_running_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint_long_running + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_long_running_with_metadata` + interceptor in new development instead of the `post_update_endpoint_long_running` interceptor. + When both interceptors are used, this `post_update_endpoint_long_running_with_metadata` interceptor runs after the + `post_update_endpoint_long_running` interceptor. The (possibly modified) response returned by + `post_update_endpoint_long_running` will be passed to + `post_update_endpoint_long_running_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2755,6 +2964,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2898,6 +3111,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3047,6 +3264,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_deploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_deploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3192,6 +3413,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3336,6 +3561,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3490,6 +3719,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_mutate_deployed_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_mutate_deployed_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3639,6 +3872,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_undeploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_undeploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3790,6 +4027,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3945,6 +4186,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_endpoint_long_running(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_endpoint_long_running_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest_asyncio.py index 77264c58b5..379c756f10 100644 --- a/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/endpoint_service/transports/rest_asyncio.py @@ -196,12 +196,35 @@ async def post_create_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_endpoint` interceptor runs + before the `post_create_endpoint_with_metadata` interceptor. """ return response + async def post_create_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_create_endpoint_with_metadata` + interceptor in new development instead of the `post_create_endpoint` interceptor. + When both interceptors are used, this `post_create_endpoint_with_metadata` interceptor runs after the + `post_create_endpoint` interceptor. The (possibly modified) response returned by + `post_create_endpoint` will be passed to + `post_create_endpoint_with_metadata`. + """ + return response, metadata + async def pre_delete_endpoint( self, request: endpoint_service.DeleteEndpointRequest, @@ -221,12 +244,35 @@ async def post_delete_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_endpoint` interceptor runs + before the `post_delete_endpoint_with_metadata` interceptor. """ return response + async def post_delete_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_endpoint` interceptor. + When both interceptors are used, this `post_delete_endpoint_with_metadata` interceptor runs after the + `post_delete_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_endpoint` will be passed to + `post_delete_endpoint_with_metadata`. + """ + return response, metadata + async def pre_deploy_model( self, request: endpoint_service.DeployModelRequest, @@ -246,12 +292,35 @@ async def post_deploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_model` interceptor runs + before the `post_deploy_model_with_metadata` interceptor. """ return response + async def post_deploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_model_with_metadata` + interceptor in new development instead of the `post_deploy_model` interceptor. + When both interceptors are used, this `post_deploy_model_with_metadata` interceptor runs after the + `post_deploy_model` interceptor. The (possibly modified) response returned by + `post_deploy_model` will be passed to + `post_deploy_model_with_metadata`. + """ + return response, metadata + async def pre_get_endpoint( self, request: endpoint_service.GetEndpointRequest, @@ -269,12 +338,35 @@ async def pre_get_endpoint( async def post_get_endpoint(self, response: endpoint.Endpoint) -> endpoint.Endpoint: """Post-rpc interceptor for get_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_endpoint` interceptor runs + before the `post_get_endpoint_with_metadata` interceptor. """ return response + async def post_get_endpoint_with_metadata( + self, + response: endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_get_endpoint_with_metadata` + interceptor in new development instead of the `post_get_endpoint` interceptor. + When both interceptors are used, this `post_get_endpoint_with_metadata` interceptor runs after the + `post_get_endpoint` interceptor. The (possibly modified) response returned by + `post_get_endpoint` will be passed to + `post_get_endpoint_with_metadata`. + """ + return response, metadata + async def pre_list_endpoints( self, request: endpoint_service.ListEndpointsRequest, @@ -294,12 +386,37 @@ async def post_list_endpoints( ) -> endpoint_service.ListEndpointsResponse: """Post-rpc interceptor for list_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_endpoints` interceptor runs + before the `post_list_endpoints_with_metadata` interceptor. """ return response + async def post_list_endpoints_with_metadata( + self, + response: endpoint_service.ListEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + endpoint_service.ListEndpointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_list_endpoints_with_metadata` + interceptor in new development instead of the `post_list_endpoints` interceptor. + When both interceptors are used, this `post_list_endpoints_with_metadata` interceptor runs after the + `post_list_endpoints` interceptor. The (possibly modified) response returned by + `post_list_endpoints` will be passed to + `post_list_endpoints_with_metadata`. + """ + return response, metadata + async def pre_mutate_deployed_model( self, request: endpoint_service.MutateDeployedModelRequest, @@ -320,12 +437,35 @@ async def post_mutate_deployed_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_model` interceptor runs + before the `post_mutate_deployed_model_with_metadata` interceptor. """ return response + async def post_mutate_deployed_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_model_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_model` interceptor. + When both interceptors are used, this `post_mutate_deployed_model_with_metadata` interceptor runs after the + `post_mutate_deployed_model` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_model` will be passed to + `post_mutate_deployed_model_with_metadata`. + """ + return response, metadata + async def pre_undeploy_model( self, request: endpoint_service.UndeployModelRequest, @@ -345,12 +485,35 @@ async def post_undeploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_model` interceptor runs + before the `post_undeploy_model_with_metadata` interceptor. """ return response + async def post_undeploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_model_with_metadata` + interceptor in new development instead of the `post_undeploy_model` interceptor. + When both interceptors are used, this `post_undeploy_model_with_metadata` interceptor runs after the + `post_undeploy_model` interceptor. The (possibly modified) response returned by + `post_undeploy_model` will be passed to + `post_undeploy_model_with_metadata`. + """ + return response, metadata + async def pre_update_endpoint( self, request: endpoint_service.UpdateEndpointRequest, @@ -370,12 +533,35 @@ async def post_update_endpoint( ) -> gca_endpoint.Endpoint: """Post-rpc interceptor for update_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint` interceptor runs + before the `post_update_endpoint_with_metadata` interceptor. """ return response + async def post_update_endpoint_with_metadata( + self, + response: gca_endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_with_metadata` + interceptor in new development instead of the `post_update_endpoint` interceptor. + When both interceptors are used, this `post_update_endpoint_with_metadata` interceptor runs after the + `post_update_endpoint` interceptor. The (possibly modified) response returned by + `post_update_endpoint` will be passed to + `post_update_endpoint_with_metadata`. + """ + return response, metadata + async def pre_update_endpoint_long_running( self, request: endpoint_service.UpdateEndpointLongRunningRequest, @@ -396,12 +582,35 @@ async def post_update_endpoint_long_running( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_endpoint_long_running - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_long_running_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint_long_running` interceptor runs + before the `post_update_endpoint_long_running_with_metadata` interceptor. """ return response + async def post_update_endpoint_long_running_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint_long_running + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_long_running_with_metadata` + interceptor in new development instead of the `post_update_endpoint_long_running` interceptor. + When both interceptors are used, this `post_update_endpoint_long_running_with_metadata` interceptor runs after the + `post_update_endpoint_long_running` interceptor. The (possibly modified) response returned by + `post_update_endpoint_long_running` will be passed to + `post_update_endpoint_long_running_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -957,6 +1166,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1113,6 +1326,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1274,6 +1491,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_deploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_deploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1429,6 +1650,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1584,6 +1809,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1746,6 +1975,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_mutate_deployed_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_mutate_deployed_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1908,6 +2141,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_undeploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_undeploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2070,6 +2307,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2234,6 +2475,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_endpoint_long_running(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_endpoint_long_running_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/evaluation_service/client.py b/google/cloud/aiplatform_v1/services/evaluation_service/client.py index 0dc5d30f5e..ce871569cd 100644 --- a/google/cloud/aiplatform_v1/services/evaluation_service/client.py +++ b/google/cloud/aiplatform_v1/services/evaluation_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -478,6 +480,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -847,16 +876,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -902,16 +935,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1074,16 +1111,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1195,16 +1236,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1317,16 +1362,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1377,16 +1426,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1432,16 +1485,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1487,16 +1544,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest.py b/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest.py index bb24ccdc1f..59dbb61f90 100644 --- a/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest.py @@ -112,12 +112,38 @@ def post_evaluate_instances( ) -> evaluation_service.EvaluateInstancesResponse: """Post-rpc interceptor for evaluate_instances - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_evaluate_instances_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EvaluationService server but before - it is returned to user code. + it is returned to user code. This `post_evaluate_instances` interceptor runs + before the `post_evaluate_instances_with_metadata` interceptor. """ return response + def post_evaluate_instances_with_metadata( + self, + response: evaluation_service.EvaluateInstancesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateInstancesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for evaluate_instances + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_instances_with_metadata` + interceptor in new development instead of the `post_evaluate_instances` interceptor. + When both interceptors are used, this `post_evaluate_instances_with_metadata` interceptor runs after the + `post_evaluate_instances` interceptor. The (possibly modified) response returned by + `post_evaluate_instances` will be passed to + `post_evaluate_instances_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -578,6 +604,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_evaluate_instances(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_evaluate_instances_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest_asyncio.py index c85eee6302..82b8b409e4 100644 --- a/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/evaluation_service/transports/rest_asyncio.py @@ -129,12 +129,38 @@ async def post_evaluate_instances( ) -> evaluation_service.EvaluateInstancesResponse: """Post-rpc interceptor for evaluate_instances - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_evaluate_instances_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EvaluationService server but before - it is returned to user code. + it is returned to user code. This `post_evaluate_instances` interceptor runs + before the `post_evaluate_instances_with_metadata` interceptor. """ return response + async def post_evaluate_instances_with_metadata( + self, + response: evaluation_service.EvaluateInstancesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateInstancesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for evaluate_instances + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_instances_with_metadata` + interceptor in new development instead of the `post_evaluate_instances` interceptor. + When both interceptors are used, this `post_evaluate_instances_with_metadata` interceptor runs after the + `post_evaluate_instances` interceptor. The (possibly modified) response returned by + `post_evaluate_instances` will be passed to + `post_evaluate_instances_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -644,6 +670,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_evaluate_instances(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_evaluate_instances_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/client.py b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/client.py index 6b72e2823d..4da4acfe65 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/client.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -582,6 +584,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2654,16 +2683,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2709,16 +2742,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2881,16 +2918,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3002,16 +3043,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3124,16 +3169,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3184,16 +3233,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3239,16 +3292,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3294,16 +3351,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest.py b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest.py index 373cf9beae..69f6c0ffde 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest.py @@ -212,12 +212,35 @@ def post_create_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_online_store` interceptor runs + before the `post_create_feature_online_store_with_metadata` interceptor. """ return response + def post_create_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_online_store_with_metadata` + interceptor in new development instead of the `post_create_feature_online_store` interceptor. + When both interceptors are used, this `post_create_feature_online_store_with_metadata` interceptor runs after the + `post_create_feature_online_store` interceptor. The (possibly modified) response returned by + `post_create_feature_online_store` will be passed to + `post_create_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_create_feature_view( self, request: feature_online_store_admin_service.CreateFeatureViewRequest, @@ -238,12 +261,35 @@ def post_create_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_view` interceptor runs + before the `post_create_feature_view_with_metadata` interceptor. """ return response + def post_create_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_view_with_metadata` + interceptor in new development instead of the `post_create_feature_view` interceptor. + When both interceptors are used, this `post_create_feature_view_with_metadata` interceptor runs after the + `post_create_feature_view` interceptor. The (possibly modified) response returned by + `post_create_feature_view` will be passed to + `post_create_feature_view_with_metadata`. + """ + return response, metadata + def pre_delete_feature_online_store( self, request: feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest, @@ -264,12 +310,35 @@ def post_delete_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_online_store` interceptor runs + before the `post_delete_feature_online_store_with_metadata` interceptor. """ return response + def post_delete_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_online_store_with_metadata` + interceptor in new development instead of the `post_delete_feature_online_store` interceptor. + When both interceptors are used, this `post_delete_feature_online_store_with_metadata` interceptor runs after the + `post_delete_feature_online_store` interceptor. The (possibly modified) response returned by + `post_delete_feature_online_store` will be passed to + `post_delete_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_delete_feature_view( self, request: feature_online_store_admin_service.DeleteFeatureViewRequest, @@ -290,12 +359,35 @@ def post_delete_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_view` interceptor runs + before the `post_delete_feature_view_with_metadata` interceptor. """ return response + def post_delete_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_view_with_metadata` + interceptor in new development instead of the `post_delete_feature_view` interceptor. + When both interceptors are used, this `post_delete_feature_view_with_metadata` interceptor runs after the + `post_delete_feature_view` interceptor. The (possibly modified) response returned by + `post_delete_feature_view` will be passed to + `post_delete_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_feature_online_store( self, request: feature_online_store_admin_service.GetFeatureOnlineStoreRequest, @@ -316,12 +408,37 @@ def post_get_feature_online_store( ) -> feature_online_store.FeatureOnlineStore: """Post-rpc interceptor for get_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_online_store` interceptor runs + before the `post_get_feature_online_store_with_metadata` interceptor. """ return response + def post_get_feature_online_store_with_metadata( + self, + response: feature_online_store.FeatureOnlineStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store.FeatureOnlineStore, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_online_store_with_metadata` + interceptor in new development instead of the `post_get_feature_online_store` interceptor. + When both interceptors are used, this `post_get_feature_online_store_with_metadata` interceptor runs after the + `post_get_feature_online_store` interceptor. The (possibly modified) response returned by + `post_get_feature_online_store` will be passed to + `post_get_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_get_feature_view( self, request: feature_online_store_admin_service.GetFeatureViewRequest, @@ -342,12 +459,35 @@ def post_get_feature_view( ) -> feature_view.FeatureView: """Post-rpc interceptor for get_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view` interceptor runs + before the `post_get_feature_view_with_metadata` interceptor. """ return response + def post_get_feature_view_with_metadata( + self, + response: feature_view.FeatureView, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_view.FeatureView, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_with_metadata` + interceptor in new development instead of the `post_get_feature_view` interceptor. + When both interceptors are used, this `post_get_feature_view_with_metadata` interceptor runs after the + `post_get_feature_view` interceptor. The (possibly modified) response returned by + `post_get_feature_view` will be passed to + `post_get_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_feature_view_sync( self, request: feature_online_store_admin_service.GetFeatureViewSyncRequest, @@ -368,12 +508,37 @@ def post_get_feature_view_sync( ) -> feature_view_sync.FeatureViewSync: """Post-rpc interceptor for get_feature_view_sync - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_sync_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view_sync` interceptor runs + before the `post_get_feature_view_sync_with_metadata` interceptor. """ return response + def post_get_feature_view_sync_with_metadata( + self, + response: feature_view_sync.FeatureViewSync, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_view_sync.FeatureViewSync, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_view_sync + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_sync_with_metadata` + interceptor in new development instead of the `post_get_feature_view_sync` interceptor. + When both interceptors are used, this `post_get_feature_view_sync_with_metadata` interceptor runs after the + `post_get_feature_view_sync` interceptor. The (possibly modified) response returned by + `post_get_feature_view_sync` will be passed to + `post_get_feature_view_sync_with_metadata`. + """ + return response, metadata + def pre_list_feature_online_stores( self, request: feature_online_store_admin_service.ListFeatureOnlineStoresRequest, @@ -395,12 +560,38 @@ def post_list_feature_online_stores( ) -> feature_online_store_admin_service.ListFeatureOnlineStoresResponse: """Post-rpc interceptor for list_feature_online_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_online_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_online_stores` interceptor runs + before the `post_list_feature_online_stores_with_metadata` interceptor. """ return response + def post_list_feature_online_stores_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_online_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_online_stores_with_metadata` + interceptor in new development instead of the `post_list_feature_online_stores` interceptor. + When both interceptors are used, this `post_list_feature_online_stores_with_metadata` interceptor runs after the + `post_list_feature_online_stores` interceptor. The (possibly modified) response returned by + `post_list_feature_online_stores` will be passed to + `post_list_feature_online_stores_with_metadata`. + """ + return response, metadata + def pre_list_feature_views( self, request: feature_online_store_admin_service.ListFeatureViewsRequest, @@ -421,12 +612,38 @@ def post_list_feature_views( ) -> feature_online_store_admin_service.ListFeatureViewsResponse: """Post-rpc interceptor for list_feature_views - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_views_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_views` interceptor runs + before the `post_list_feature_views_with_metadata` interceptor. """ return response + def post_list_feature_views_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_views + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_views_with_metadata` + interceptor in new development instead of the `post_list_feature_views` interceptor. + When both interceptors are used, this `post_list_feature_views_with_metadata` interceptor runs after the + `post_list_feature_views` interceptor. The (possibly modified) response returned by + `post_list_feature_views` will be passed to + `post_list_feature_views_with_metadata`. + """ + return response, metadata + def pre_list_feature_view_syncs( self, request: feature_online_store_admin_service.ListFeatureViewSyncsRequest, @@ -447,12 +664,38 @@ def post_list_feature_view_syncs( ) -> feature_online_store_admin_service.ListFeatureViewSyncsResponse: """Post-rpc interceptor for list_feature_view_syncs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_view_syncs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_view_syncs` interceptor runs + before the `post_list_feature_view_syncs_with_metadata` interceptor. """ return response + def post_list_feature_view_syncs_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewSyncsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewSyncsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_view_syncs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_view_syncs_with_metadata` + interceptor in new development instead of the `post_list_feature_view_syncs` interceptor. + When both interceptors are used, this `post_list_feature_view_syncs_with_metadata` interceptor runs after the + `post_list_feature_view_syncs` interceptor. The (possibly modified) response returned by + `post_list_feature_view_syncs` will be passed to + `post_list_feature_view_syncs_with_metadata`. + """ + return response, metadata + def pre_sync_feature_view( self, request: feature_online_store_admin_service.SyncFeatureViewRequest, @@ -473,12 +716,38 @@ def post_sync_feature_view( ) -> feature_online_store_admin_service.SyncFeatureViewResponse: """Post-rpc interceptor for sync_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_sync_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_sync_feature_view` interceptor runs + before the `post_sync_feature_view_with_metadata` interceptor. """ return response + def post_sync_feature_view_with_metadata( + self, + response: feature_online_store_admin_service.SyncFeatureViewResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.SyncFeatureViewResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for sync_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_sync_feature_view_with_metadata` + interceptor in new development instead of the `post_sync_feature_view` interceptor. + When both interceptors are used, this `post_sync_feature_view_with_metadata` interceptor runs after the + `post_sync_feature_view` interceptor. The (possibly modified) response returned by + `post_sync_feature_view` will be passed to + `post_sync_feature_view_with_metadata`. + """ + return response, metadata + def pre_update_feature_online_store( self, request: feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest, @@ -499,12 +768,35 @@ def post_update_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_online_store` interceptor runs + before the `post_update_feature_online_store_with_metadata` interceptor. """ return response + def post_update_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_online_store_with_metadata` + interceptor in new development instead of the `post_update_feature_online_store` interceptor. + When both interceptors are used, this `post_update_feature_online_store_with_metadata` interceptor runs after the + `post_update_feature_online_store` interceptor. The (possibly modified) response returned by + `post_update_feature_online_store` will be passed to + `post_update_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_update_feature_view( self, request: feature_online_store_admin_service.UpdateFeatureViewRequest, @@ -525,12 +817,35 @@ def post_update_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_view` interceptor runs + before the `post_update_feature_view_with_metadata` interceptor. """ return response + def post_update_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_view_with_metadata` + interceptor in new development instead of the `post_update_feature_view` interceptor. + When both interceptors are used, this `post_update_feature_view_with_metadata` interceptor runs after the + `post_update_feature_view` interceptor. The (possibly modified) response returned by + `post_update_feature_view` will be passed to + `post_update_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2913,6 +3228,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3065,6 +3384,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3214,6 +3537,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3360,6 +3687,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3512,6 +3843,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3662,6 +3997,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3813,6 +4152,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_view_sync(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_view_sync_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3969,6 +4312,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_online_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_online_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4120,6 +4467,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_views(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_views_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4273,6 +4624,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_view_syncs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_view_syncs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4430,6 +4785,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_sync_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_sync_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4587,6 +4946,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4739,6 +5102,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest_asyncio.py index d0139b4521..3ddcc171be 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_admin_service/transports/rest_asyncio.py @@ -230,12 +230,35 @@ async def post_create_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_online_store` interceptor runs + before the `post_create_feature_online_store_with_metadata` interceptor. """ return response + async def post_create_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_online_store_with_metadata` + interceptor in new development instead of the `post_create_feature_online_store` interceptor. + When both interceptors are used, this `post_create_feature_online_store_with_metadata` interceptor runs after the + `post_create_feature_online_store` interceptor. The (possibly modified) response returned by + `post_create_feature_online_store` will be passed to + `post_create_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_create_feature_view( self, request: feature_online_store_admin_service.CreateFeatureViewRequest, @@ -256,12 +279,35 @@ async def post_create_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_view` interceptor runs + before the `post_create_feature_view_with_metadata` interceptor. """ return response + async def post_create_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_view_with_metadata` + interceptor in new development instead of the `post_create_feature_view` interceptor. + When both interceptors are used, this `post_create_feature_view_with_metadata` interceptor runs after the + `post_create_feature_view` interceptor. The (possibly modified) response returned by + `post_create_feature_view` will be passed to + `post_create_feature_view_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_online_store( self, request: feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest, @@ -282,12 +328,35 @@ async def post_delete_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_online_store` interceptor runs + before the `post_delete_feature_online_store_with_metadata` interceptor. """ return response + async def post_delete_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_online_store_with_metadata` + interceptor in new development instead of the `post_delete_feature_online_store` interceptor. + When both interceptors are used, this `post_delete_feature_online_store_with_metadata` interceptor runs after the + `post_delete_feature_online_store` interceptor. The (possibly modified) response returned by + `post_delete_feature_online_store` will be passed to + `post_delete_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_view( self, request: feature_online_store_admin_service.DeleteFeatureViewRequest, @@ -308,12 +377,35 @@ async def post_delete_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_view` interceptor runs + before the `post_delete_feature_view_with_metadata` interceptor. """ return response + async def post_delete_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_view_with_metadata` + interceptor in new development instead of the `post_delete_feature_view` interceptor. + When both interceptors are used, this `post_delete_feature_view_with_metadata` interceptor runs after the + `post_delete_feature_view` interceptor. The (possibly modified) response returned by + `post_delete_feature_view` will be passed to + `post_delete_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_feature_online_store( self, request: feature_online_store_admin_service.GetFeatureOnlineStoreRequest, @@ -334,12 +426,37 @@ async def post_get_feature_online_store( ) -> feature_online_store.FeatureOnlineStore: """Post-rpc interceptor for get_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_online_store` interceptor runs + before the `post_get_feature_online_store_with_metadata` interceptor. """ return response + async def post_get_feature_online_store_with_metadata( + self, + response: feature_online_store.FeatureOnlineStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store.FeatureOnlineStore, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_online_store_with_metadata` + interceptor in new development instead of the `post_get_feature_online_store` interceptor. + When both interceptors are used, this `post_get_feature_online_store_with_metadata` interceptor runs after the + `post_get_feature_online_store` interceptor. The (possibly modified) response returned by + `post_get_feature_online_store` will be passed to + `post_get_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_get_feature_view( self, request: feature_online_store_admin_service.GetFeatureViewRequest, @@ -360,12 +477,35 @@ async def post_get_feature_view( ) -> feature_view.FeatureView: """Post-rpc interceptor for get_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view` interceptor runs + before the `post_get_feature_view_with_metadata` interceptor. """ return response + async def post_get_feature_view_with_metadata( + self, + response: feature_view.FeatureView, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_view.FeatureView, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_with_metadata` + interceptor in new development instead of the `post_get_feature_view` interceptor. + When both interceptors are used, this `post_get_feature_view_with_metadata` interceptor runs after the + `post_get_feature_view` interceptor. The (possibly modified) response returned by + `post_get_feature_view` will be passed to + `post_get_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_feature_view_sync( self, request: feature_online_store_admin_service.GetFeatureViewSyncRequest, @@ -386,12 +526,37 @@ async def post_get_feature_view_sync( ) -> feature_view_sync.FeatureViewSync: """Post-rpc interceptor for get_feature_view_sync - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_sync_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view_sync` interceptor runs + before the `post_get_feature_view_sync_with_metadata` interceptor. """ return response + async def post_get_feature_view_sync_with_metadata( + self, + response: feature_view_sync.FeatureViewSync, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_view_sync.FeatureViewSync, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_view_sync + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_sync_with_metadata` + interceptor in new development instead of the `post_get_feature_view_sync` interceptor. + When both interceptors are used, this `post_get_feature_view_sync_with_metadata` interceptor runs after the + `post_get_feature_view_sync` interceptor. The (possibly modified) response returned by + `post_get_feature_view_sync` will be passed to + `post_get_feature_view_sync_with_metadata`. + """ + return response, metadata + async def pre_list_feature_online_stores( self, request: feature_online_store_admin_service.ListFeatureOnlineStoresRequest, @@ -413,12 +578,38 @@ async def post_list_feature_online_stores( ) -> feature_online_store_admin_service.ListFeatureOnlineStoresResponse: """Post-rpc interceptor for list_feature_online_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_online_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_online_stores` interceptor runs + before the `post_list_feature_online_stores_with_metadata` interceptor. """ return response + async def post_list_feature_online_stores_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_online_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_online_stores_with_metadata` + interceptor in new development instead of the `post_list_feature_online_stores` interceptor. + When both interceptors are used, this `post_list_feature_online_stores_with_metadata` interceptor runs after the + `post_list_feature_online_stores` interceptor. The (possibly modified) response returned by + `post_list_feature_online_stores` will be passed to + `post_list_feature_online_stores_with_metadata`. + """ + return response, metadata + async def pre_list_feature_views( self, request: feature_online_store_admin_service.ListFeatureViewsRequest, @@ -439,12 +630,38 @@ async def post_list_feature_views( ) -> feature_online_store_admin_service.ListFeatureViewsResponse: """Post-rpc interceptor for list_feature_views - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_views_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_views` interceptor runs + before the `post_list_feature_views_with_metadata` interceptor. """ return response + async def post_list_feature_views_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_views + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_views_with_metadata` + interceptor in new development instead of the `post_list_feature_views` interceptor. + When both interceptors are used, this `post_list_feature_views_with_metadata` interceptor runs after the + `post_list_feature_views` interceptor. The (possibly modified) response returned by + `post_list_feature_views` will be passed to + `post_list_feature_views_with_metadata`. + """ + return response, metadata + async def pre_list_feature_view_syncs( self, request: feature_online_store_admin_service.ListFeatureViewSyncsRequest, @@ -465,12 +682,38 @@ async def post_list_feature_view_syncs( ) -> feature_online_store_admin_service.ListFeatureViewSyncsResponse: """Post-rpc interceptor for list_feature_view_syncs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_view_syncs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_view_syncs` interceptor runs + before the `post_list_feature_view_syncs_with_metadata` interceptor. """ return response + async def post_list_feature_view_syncs_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewSyncsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewSyncsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_view_syncs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_view_syncs_with_metadata` + interceptor in new development instead of the `post_list_feature_view_syncs` interceptor. + When both interceptors are used, this `post_list_feature_view_syncs_with_metadata` interceptor runs after the + `post_list_feature_view_syncs` interceptor. The (possibly modified) response returned by + `post_list_feature_view_syncs` will be passed to + `post_list_feature_view_syncs_with_metadata`. + """ + return response, metadata + async def pre_sync_feature_view( self, request: feature_online_store_admin_service.SyncFeatureViewRequest, @@ -491,12 +734,38 @@ async def post_sync_feature_view( ) -> feature_online_store_admin_service.SyncFeatureViewResponse: """Post-rpc interceptor for sync_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_sync_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_sync_feature_view` interceptor runs + before the `post_sync_feature_view_with_metadata` interceptor. """ return response + async def post_sync_feature_view_with_metadata( + self, + response: feature_online_store_admin_service.SyncFeatureViewResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.SyncFeatureViewResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for sync_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_sync_feature_view_with_metadata` + interceptor in new development instead of the `post_sync_feature_view` interceptor. + When both interceptors are used, this `post_sync_feature_view_with_metadata` interceptor runs after the + `post_sync_feature_view` interceptor. The (possibly modified) response returned by + `post_sync_feature_view` will be passed to + `post_sync_feature_view_with_metadata`. + """ + return response, metadata + async def pre_update_feature_online_store( self, request: feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest, @@ -517,12 +786,35 @@ async def post_update_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_online_store` interceptor runs + before the `post_update_feature_online_store_with_metadata` interceptor. """ return response + async def post_update_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_online_store_with_metadata` + interceptor in new development instead of the `post_update_feature_online_store` interceptor. + When both interceptors are used, this `post_update_feature_online_store_with_metadata` interceptor runs after the + `post_update_feature_online_store` interceptor. The (possibly modified) response returned by + `post_update_feature_online_store` will be passed to + `post_update_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_update_feature_view( self, request: feature_online_store_admin_service.UpdateFeatureViewRequest, @@ -543,12 +835,35 @@ async def post_update_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_view` interceptor runs + before the `post_update_feature_view_with_metadata` interceptor. """ return response + async def post_update_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_view_with_metadata` + interceptor in new development instead of the `post_update_feature_view` interceptor. + When both interceptors are used, this `post_update_feature_view_with_metadata` interceptor runs after the + `post_update_feature_view` interceptor. The (possibly modified) response returned by + `post_update_feature_view` will be passed to + `post_update_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1132,6 +1447,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1294,6 +1616,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1451,6 +1777,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1607,6 +1940,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1765,6 +2102,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1923,6 +2267,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2080,6 +2428,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_view_sync(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_view_sync_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2242,6 +2594,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_online_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_feature_online_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2401,6 +2760,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_views(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_feature_views_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2560,6 +2923,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_view_syncs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_feature_view_syncs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2725,6 +3095,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_sync_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_sync_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2890,6 +3264,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3052,6 +3433,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_service/client.py b/google/cloud/aiplatform_v1/services/feature_online_store_service/client.py index 1484975278..ca08f2f82a 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_service/client.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -507,6 +509,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1012,16 +1041,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1067,16 +1100,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1239,16 +1276,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1360,16 +1401,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1482,16 +1527,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1542,16 +1591,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1597,16 +1650,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1652,16 +1709,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest.py b/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest.py index 5fea64417c..0fa86f8a0d 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest.py @@ -120,12 +120,38 @@ def post_fetch_feature_values( ) -> feature_online_store_service.FetchFeatureValuesResponse: """Post-rpc interceptor for fetch_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_fetch_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_fetch_feature_values` interceptor runs + before the `post_fetch_feature_values_with_metadata` interceptor. """ return response + def post_fetch_feature_values_with_metadata( + self, + response: feature_online_store_service.FetchFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.FetchFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for fetch_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_fetch_feature_values_with_metadata` + interceptor in new development instead of the `post_fetch_feature_values` interceptor. + When both interceptors are used, this `post_fetch_feature_values_with_metadata` interceptor runs after the + `post_fetch_feature_values` interceptor. The (possibly modified) response returned by + `post_fetch_feature_values` will be passed to + `post_fetch_feature_values_with_metadata`. + """ + return response, metadata + def pre_search_nearest_entities( self, request: feature_online_store_service.SearchNearestEntitiesRequest, @@ -146,12 +172,38 @@ def post_search_nearest_entities( ) -> feature_online_store_service.SearchNearestEntitiesResponse: """Post-rpc interceptor for search_nearest_entities - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_nearest_entities_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_nearest_entities` interceptor runs + before the `post_search_nearest_entities_with_metadata` interceptor. """ return response + def post_search_nearest_entities_with_metadata( + self, + response: feature_online_store_service.SearchNearestEntitiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.SearchNearestEntitiesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_nearest_entities + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_search_nearest_entities_with_metadata` + interceptor in new development instead of the `post_search_nearest_entities` interceptor. + When both interceptors are used, this `post_search_nearest_entities_with_metadata` interceptor runs after the + `post_search_nearest_entities` interceptor. The (possibly modified) response returned by + `post_search_nearest_entities` will be passed to + `post_search_nearest_entities_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -616,6 +668,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_fetch_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_fetch_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -775,6 +831,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_nearest_entities(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_nearest_entities_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest_asyncio.py index e13eceebc4..c0d3c1022e 100644 --- a/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/feature_online_store_service/transports/rest_asyncio.py @@ -137,12 +137,38 @@ async def post_fetch_feature_values( ) -> feature_online_store_service.FetchFeatureValuesResponse: """Post-rpc interceptor for fetch_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_fetch_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_fetch_feature_values` interceptor runs + before the `post_fetch_feature_values_with_metadata` interceptor. """ return response + async def post_fetch_feature_values_with_metadata( + self, + response: feature_online_store_service.FetchFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.FetchFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for fetch_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_fetch_feature_values_with_metadata` + interceptor in new development instead of the `post_fetch_feature_values` interceptor. + When both interceptors are used, this `post_fetch_feature_values_with_metadata` interceptor runs after the + `post_fetch_feature_values` interceptor. The (possibly modified) response returned by + `post_fetch_feature_values` will be passed to + `post_fetch_feature_values_with_metadata`. + """ + return response, metadata + async def pre_search_nearest_entities( self, request: feature_online_store_service.SearchNearestEntitiesRequest, @@ -163,12 +189,38 @@ async def post_search_nearest_entities( ) -> feature_online_store_service.SearchNearestEntitiesResponse: """Post-rpc interceptor for search_nearest_entities - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_nearest_entities_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_nearest_entities` interceptor runs + before the `post_search_nearest_entities_with_metadata` interceptor. """ return response + async def post_search_nearest_entities_with_metadata( + self, + response: feature_online_store_service.SearchNearestEntitiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.SearchNearestEntitiesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_nearest_entities + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_search_nearest_entities_with_metadata` + interceptor in new development instead of the `post_search_nearest_entities` interceptor. + When both interceptors are used, this `post_search_nearest_entities_with_metadata` interceptor runs after the + `post_search_nearest_entities` interceptor. The (possibly modified) response returned by + `post_search_nearest_entities` will be passed to + `post_search_nearest_entities_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -691,6 +743,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_fetch_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_fetch_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -858,6 +914,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_nearest_entities(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_nearest_entities_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_registry_service/client.py b/google/cloud/aiplatform_v1/services/feature_registry_service/client.py index 7ded4d3d77..370e79a9f2 100644 --- a/google/cloud/aiplatform_v1/services/feature_registry_service/client.py +++ b/google/cloud/aiplatform_v1/services/feature_registry_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -542,6 +544,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2326,16 +2355,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2381,16 +2414,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2553,16 +2590,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2674,16 +2715,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2796,16 +2841,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2856,16 +2905,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2911,16 +2964,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2966,16 +3023,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest.py b/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest.py index beab4bba51..f4ea1e44a0 100644 --- a/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest.py @@ -196,12 +196,35 @@ def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -222,12 +245,35 @@ def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + def pre_create_feature_group( self, request: feature_registry_service.CreateFeatureGroupRequest, @@ -248,12 +294,35 @@ def post_create_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_group` interceptor runs + before the `post_create_feature_group_with_metadata` interceptor. """ return response + def post_create_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_group_with_metadata` + interceptor in new development instead of the `post_create_feature_group` interceptor. + When both interceptors are used, this `post_create_feature_group_with_metadata` interceptor runs after the + `post_create_feature_group` interceptor. The (possibly modified) response returned by + `post_create_feature_group` will be passed to + `post_create_feature_group_with_metadata`. + """ + return response, metadata + def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -274,12 +343,35 @@ def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + def pre_delete_feature_group( self, request: feature_registry_service.DeleteFeatureGroupRequest, @@ -300,12 +392,35 @@ def post_delete_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_group` interceptor runs + before the `post_delete_feature_group_with_metadata` interceptor. """ return response + def post_delete_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_group_with_metadata` + interceptor in new development instead of the `post_delete_feature_group` interceptor. + When both interceptors are used, this `post_delete_feature_group_with_metadata` interceptor runs after the + `post_delete_feature_group` interceptor. The (possibly modified) response returned by + `post_delete_feature_group` will be passed to + `post_delete_feature_group_with_metadata`. + """ + return response, metadata + def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -323,12 +438,35 @@ def pre_get_feature( def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + def pre_get_feature_group( self, request: feature_registry_service.GetFeatureGroupRequest, @@ -349,12 +487,35 @@ def post_get_feature_group( ) -> feature_group.FeatureGroup: """Post-rpc interceptor for get_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_group` interceptor runs + before the `post_get_feature_group_with_metadata` interceptor. """ return response + def post_get_feature_group_with_metadata( + self, + response: feature_group.FeatureGroup, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_group.FeatureGroup, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_group_with_metadata` + interceptor in new development instead of the `post_get_feature_group` interceptor. + When both interceptors are used, this `post_get_feature_group_with_metadata` interceptor runs after the + `post_get_feature_group` interceptor. The (possibly modified) response returned by + `post_get_feature_group` will be passed to + `post_get_feature_group_with_metadata`. + """ + return response, metadata + def pre_list_feature_groups( self, request: feature_registry_service.ListFeatureGroupsRequest, @@ -375,12 +536,38 @@ def post_list_feature_groups( ) -> feature_registry_service.ListFeatureGroupsResponse: """Post-rpc interceptor for list_feature_groups - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_groups_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_groups` interceptor runs + before the `post_list_feature_groups_with_metadata` interceptor. """ return response + def post_list_feature_groups_with_metadata( + self, + response: feature_registry_service.ListFeatureGroupsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureGroupsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_groups + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_groups_with_metadata` + interceptor in new development instead of the `post_list_feature_groups` interceptor. + When both interceptors are used, this `post_list_feature_groups_with_metadata` interceptor runs after the + `post_list_feature_groups` interceptor. The (possibly modified) response returned by + `post_list_feature_groups` will be passed to + `post_list_feature_groups_with_metadata`. + """ + return response, metadata + def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -401,12 +588,38 @@ def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -427,12 +640,35 @@ def post_update_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + def post_update_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + def pre_update_feature_group( self, request: feature_registry_service.UpdateFeatureGroupRequest, @@ -453,12 +689,35 @@ def post_update_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_group` interceptor runs + before the `post_update_feature_group_with_metadata` interceptor. """ return response + def post_update_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_group_with_metadata` + interceptor in new development instead of the `post_update_feature_group` interceptor. + When both interceptors are used, this `post_update_feature_group_with_metadata` interceptor runs after the + `post_update_feature_group` interceptor. The (possibly modified) response returned by + `post_update_feature_group` will be passed to + `post_update_feature_group_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2838,6 +3097,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2990,6 +3253,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3144,6 +3411,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3290,6 +3561,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3438,6 +3713,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3586,6 +3865,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3733,6 +4016,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3882,6 +4169,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_groups(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_groups_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4035,6 +4326,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4189,6 +4484,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4343,6 +4642,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest_asyncio.py index f278ecebe5..f5e323f301 100644 --- a/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/feature_registry_service/transports/rest_asyncio.py @@ -214,12 +214,35 @@ async def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + async def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + async def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -240,12 +263,35 @@ async def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + async def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + async def pre_create_feature_group( self, request: feature_registry_service.CreateFeatureGroupRequest, @@ -266,12 +312,35 @@ async def post_create_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_group` interceptor runs + before the `post_create_feature_group_with_metadata` interceptor. """ return response + async def post_create_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_group_with_metadata` + interceptor in new development instead of the `post_create_feature_group` interceptor. + When both interceptors are used, this `post_create_feature_group_with_metadata` interceptor runs after the + `post_create_feature_group` interceptor. The (possibly modified) response returned by + `post_create_feature_group` will be passed to + `post_create_feature_group_with_metadata`. + """ + return response, metadata + async def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -292,12 +361,35 @@ async def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + async def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_group( self, request: feature_registry_service.DeleteFeatureGroupRequest, @@ -318,12 +410,35 @@ async def post_delete_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_group` interceptor runs + before the `post_delete_feature_group_with_metadata` interceptor. """ return response + async def post_delete_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_group_with_metadata` + interceptor in new development instead of the `post_delete_feature_group` interceptor. + When both interceptors are used, this `post_delete_feature_group_with_metadata` interceptor runs after the + `post_delete_feature_group` interceptor. The (possibly modified) response returned by + `post_delete_feature_group` will be passed to + `post_delete_feature_group_with_metadata`. + """ + return response, metadata + async def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -341,12 +456,35 @@ async def pre_get_feature( async def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + async def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + async def pre_get_feature_group( self, request: feature_registry_service.GetFeatureGroupRequest, @@ -367,12 +505,35 @@ async def post_get_feature_group( ) -> feature_group.FeatureGroup: """Post-rpc interceptor for get_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_group` interceptor runs + before the `post_get_feature_group_with_metadata` interceptor. """ return response + async def post_get_feature_group_with_metadata( + self, + response: feature_group.FeatureGroup, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_group.FeatureGroup, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_group_with_metadata` + interceptor in new development instead of the `post_get_feature_group` interceptor. + When both interceptors are used, this `post_get_feature_group_with_metadata` interceptor runs after the + `post_get_feature_group` interceptor. The (possibly modified) response returned by + `post_get_feature_group` will be passed to + `post_get_feature_group_with_metadata`. + """ + return response, metadata + async def pre_list_feature_groups( self, request: feature_registry_service.ListFeatureGroupsRequest, @@ -393,12 +554,38 @@ async def post_list_feature_groups( ) -> feature_registry_service.ListFeatureGroupsResponse: """Post-rpc interceptor for list_feature_groups - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_groups_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_groups` interceptor runs + before the `post_list_feature_groups_with_metadata` interceptor. """ return response + async def post_list_feature_groups_with_metadata( + self, + response: feature_registry_service.ListFeatureGroupsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureGroupsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_groups + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_groups_with_metadata` + interceptor in new development instead of the `post_list_feature_groups` interceptor. + When both interceptors are used, this `post_list_feature_groups_with_metadata` interceptor runs after the + `post_list_feature_groups` interceptor. The (possibly modified) response returned by + `post_list_feature_groups` will be passed to + `post_list_feature_groups_with_metadata`. + """ + return response, metadata + async def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -419,12 +606,38 @@ async def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + async def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + async def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -445,12 +658,35 @@ async def post_update_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + async def post_update_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + async def pre_update_feature_group( self, request: feature_registry_service.UpdateFeatureGroupRequest, @@ -471,12 +707,35 @@ async def post_update_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_group` interceptor runs + before the `post_update_feature_group_with_metadata` interceptor. """ return response + async def post_update_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_group_with_metadata` + interceptor in new development instead of the `post_update_feature_group` interceptor. + When both interceptors are used, this `post_update_feature_group_with_metadata` interceptor runs after the + `post_update_feature_group` interceptor. The (possibly modified) response returned by + `post_update_feature_group` will be passed to + `post_update_feature_group_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1045,6 +1304,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1207,6 +1470,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1367,6 +1634,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1523,6 +1794,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1677,6 +1952,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1833,6 +2112,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1984,6 +2267,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2137,6 +2424,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_groups(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_feature_groups_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2298,6 +2589,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2462,6 +2757,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2622,6 +2921,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/client.py b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/client.py index 4ad99e8df7..b3baf44585 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/client.py +++ b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -516,6 +518,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1192,16 +1221,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1247,16 +1280,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1419,16 +1456,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1540,16 +1581,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1662,16 +1707,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1722,16 +1771,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1777,16 +1830,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1832,16 +1889,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest.py b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest.py index 250186b13a..b13a73e6a0 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest.py @@ -128,12 +128,38 @@ def post_read_feature_values( ) -> featurestore_online_service.ReadFeatureValuesResponse: """Post-rpc interceptor for read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_read_feature_values` interceptor runs + before the `post_read_feature_values_with_metadata` interceptor. """ return response + def post_read_feature_values_with_metadata( + self, + response: featurestore_online_service.ReadFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.ReadFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_read_feature_values_with_metadata` + interceptor in new development instead of the `post_read_feature_values` interceptor. + When both interceptors are used, this `post_read_feature_values_with_metadata` interceptor runs after the + `post_read_feature_values` interceptor. The (possibly modified) response returned by + `post_read_feature_values` will be passed to + `post_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_streaming_read_feature_values( self, request: featurestore_online_service.StreamingReadFeatureValuesRequest, @@ -154,12 +180,37 @@ def post_streaming_read_feature_values( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for streaming_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_streaming_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_streaming_read_feature_values` interceptor runs + before the `post_streaming_read_feature_values_with_metadata` interceptor. """ return response + def post_streaming_read_feature_values_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for streaming_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_streaming_read_feature_values_with_metadata` + interceptor in new development instead of the `post_streaming_read_feature_values` interceptor. + When both interceptors are used, this `post_streaming_read_feature_values_with_metadata` interceptor runs after the + `post_streaming_read_feature_values` interceptor. The (possibly modified) response returned by + `post_streaming_read_feature_values` will be passed to + `post_streaming_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_write_feature_values( self, request: featurestore_online_service.WriteFeatureValuesRequest, @@ -180,12 +231,38 @@ def post_write_feature_values( ) -> featurestore_online_service.WriteFeatureValuesResponse: """Post-rpc interceptor for write_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_write_feature_values` interceptor runs + before the `post_write_feature_values_with_metadata` interceptor. """ return response + def post_write_feature_values_with_metadata( + self, + response: featurestore_online_service.WriteFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.WriteFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_write_feature_values_with_metadata` + interceptor in new development instead of the `post_write_feature_values` interceptor. + When both interceptors are used, this `post_write_feature_values_with_metadata` interceptor runs after the + `post_write_feature_values` interceptor. The (possibly modified) response returned by + `post_write_feature_values` will be passed to + `post_write_feature_values_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -652,6 +729,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -812,6 +893,13 @@ def __call__( ) resp = self._interceptor.post_streaming_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_streaming_read_feature_values_with_metadata( + resp, response_metadata + ) return resp class _WriteFeatureValues( @@ -946,6 +1034,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_write_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest_asyncio.py index 0d580bd02a..1fd42912e9 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/featurestore_online_serving_service/transports/rest_asyncio.py @@ -145,12 +145,38 @@ async def post_read_feature_values( ) -> featurestore_online_service.ReadFeatureValuesResponse: """Post-rpc interceptor for read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_read_feature_values` interceptor runs + before the `post_read_feature_values_with_metadata` interceptor. """ return response + async def post_read_feature_values_with_metadata( + self, + response: featurestore_online_service.ReadFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.ReadFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_read_feature_values_with_metadata` + interceptor in new development instead of the `post_read_feature_values` interceptor. + When both interceptors are used, this `post_read_feature_values_with_metadata` interceptor runs after the + `post_read_feature_values` interceptor. The (possibly modified) response returned by + `post_read_feature_values` will be passed to + `post_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_streaming_read_feature_values( self, request: featurestore_online_service.StreamingReadFeatureValuesRequest, @@ -171,12 +197,38 @@ async def post_streaming_read_feature_values( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for streaming_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_streaming_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_streaming_read_feature_values` interceptor runs + before the `post_streaming_read_feature_values_with_metadata` interceptor. """ return response + async def post_streaming_read_feature_values_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for streaming_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_streaming_read_feature_values_with_metadata` + interceptor in new development instead of the `post_streaming_read_feature_values` interceptor. + When both interceptors are used, this `post_streaming_read_feature_values_with_metadata` interceptor runs after the + `post_streaming_read_feature_values` interceptor. The (possibly modified) response returned by + `post_streaming_read_feature_values` will be passed to + `post_streaming_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_write_feature_values( self, request: featurestore_online_service.WriteFeatureValuesRequest, @@ -197,12 +249,38 @@ async def post_write_feature_values( ) -> featurestore_online_service.WriteFeatureValuesResponse: """Post-rpc interceptor for write_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_write_feature_values` interceptor runs + before the `post_write_feature_values_with_metadata` interceptor. """ return response + async def post_write_feature_values_with_metadata( + self, + response: featurestore_online_service.WriteFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.WriteFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_write_feature_values_with_metadata` + interceptor in new development instead of the `post_write_feature_values` interceptor. + When both interceptors are used, this `post_write_feature_values_with_metadata` interceptor runs after the + `post_write_feature_values` interceptor. The (possibly modified) response returned by + `post_write_feature_values` will be passed to + `post_write_feature_values_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -730,6 +808,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -898,6 +980,13 @@ async def __call__( response, featurestore_online_service.ReadFeatureValuesResponse ) resp = await self._interceptor.post_streaming_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_streaming_read_feature_values_with_metadata( + resp, response_metadata + ) return resp class _WriteFeatureValues( @@ -1037,6 +1126,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_write_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/featurestore_service/client.py b/google/cloud/aiplatform_v1/services/featurestore_service/client.py index ef45471379..f5af8bfcb8 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_service/client.py +++ b/google/cloud/aiplatform_v1/services/featurestore_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -567,6 +569,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3788,16 +3817,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3843,16 +3876,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -4015,16 +4052,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -4136,16 +4177,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -4258,16 +4303,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -4318,16 +4367,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -4373,16 +4426,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -4428,16 +4485,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest.py b/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest.py index 06eb04a7df..d2e67ae4b4 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest.py @@ -278,12 +278,35 @@ def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + def pre_batch_read_feature_values( self, request: featurestore_service.BatchReadFeatureValuesRequest, @@ -304,12 +327,35 @@ def post_batch_read_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_feature_values` interceptor runs + before the `post_batch_read_feature_values_with_metadata` interceptor. """ return response + def post_batch_read_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_read_feature_values_with_metadata` + interceptor in new development instead of the `post_batch_read_feature_values` interceptor. + When both interceptors are used, this `post_batch_read_feature_values_with_metadata` interceptor runs after the + `post_batch_read_feature_values` interceptor. The (possibly modified) response returned by + `post_batch_read_feature_values` will be passed to + `post_batch_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_create_entity_type( self, request: featurestore_service.CreateEntityTypeRequest, @@ -330,12 +376,35 @@ def post_create_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_entity_type` interceptor runs + before the `post_create_entity_type_with_metadata` interceptor. """ return response + def post_create_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_entity_type_with_metadata` + interceptor in new development instead of the `post_create_entity_type` interceptor. + When both interceptors are used, this `post_create_entity_type_with_metadata` interceptor runs after the + `post_create_entity_type` interceptor. The (possibly modified) response returned by + `post_create_entity_type` will be passed to + `post_create_entity_type_with_metadata`. + """ + return response, metadata + def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -356,12 +425,35 @@ def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + def pre_create_featurestore( self, request: featurestore_service.CreateFeaturestoreRequest, @@ -382,12 +474,35 @@ def post_create_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_featurestore` interceptor runs + before the `post_create_featurestore_with_metadata` interceptor. """ return response + def post_create_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_featurestore_with_metadata` + interceptor in new development instead of the `post_create_featurestore` interceptor. + When both interceptors are used, this `post_create_featurestore_with_metadata` interceptor runs after the + `post_create_featurestore` interceptor. The (possibly modified) response returned by + `post_create_featurestore` will be passed to + `post_create_featurestore_with_metadata`. + """ + return response, metadata + def pre_delete_entity_type( self, request: featurestore_service.DeleteEntityTypeRequest, @@ -408,12 +523,35 @@ def post_delete_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_entity_type` interceptor runs + before the `post_delete_entity_type_with_metadata` interceptor. """ return response + def post_delete_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_entity_type_with_metadata` + interceptor in new development instead of the `post_delete_entity_type` interceptor. + When both interceptors are used, this `post_delete_entity_type_with_metadata` interceptor runs after the + `post_delete_entity_type` interceptor. The (possibly modified) response returned by + `post_delete_entity_type` will be passed to + `post_delete_entity_type_with_metadata`. + """ + return response, metadata + def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -434,12 +572,35 @@ def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + def pre_delete_featurestore( self, request: featurestore_service.DeleteFeaturestoreRequest, @@ -460,12 +621,35 @@ def post_delete_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_featurestore` interceptor runs + before the `post_delete_featurestore_with_metadata` interceptor. """ return response + def post_delete_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_featurestore_with_metadata` + interceptor in new development instead of the `post_delete_featurestore` interceptor. + When both interceptors are used, this `post_delete_featurestore_with_metadata` interceptor runs after the + `post_delete_featurestore` interceptor. The (possibly modified) response returned by + `post_delete_featurestore` will be passed to + `post_delete_featurestore_with_metadata`. + """ + return response, metadata + def pre_delete_feature_values( self, request: featurestore_service.DeleteFeatureValuesRequest, @@ -486,12 +670,35 @@ def post_delete_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_values` interceptor runs + before the `post_delete_feature_values_with_metadata` interceptor. """ return response + def post_delete_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_values_with_metadata` + interceptor in new development instead of the `post_delete_feature_values` interceptor. + When both interceptors are used, this `post_delete_feature_values_with_metadata` interceptor runs after the + `post_delete_feature_values` interceptor. The (possibly modified) response returned by + `post_delete_feature_values` will be passed to + `post_delete_feature_values_with_metadata`. + """ + return response, metadata + def pre_export_feature_values( self, request: featurestore_service.ExportFeatureValuesRequest, @@ -512,12 +719,35 @@ def post_export_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_export_feature_values` interceptor runs + before the `post_export_feature_values_with_metadata` interceptor. """ return response + def post_export_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_export_feature_values_with_metadata` + interceptor in new development instead of the `post_export_feature_values` interceptor. + When both interceptors are used, this `post_export_feature_values_with_metadata` interceptor runs after the + `post_export_feature_values` interceptor. The (possibly modified) response returned by + `post_export_feature_values` will be passed to + `post_export_feature_values_with_metadata`. + """ + return response, metadata + def pre_get_entity_type( self, request: featurestore_service.GetEntityTypeRequest, @@ -538,12 +768,35 @@ def post_get_entity_type( ) -> entity_type.EntityType: """Post-rpc interceptor for get_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_entity_type` interceptor runs + before the `post_get_entity_type_with_metadata` interceptor. """ return response + def post_get_entity_type_with_metadata( + self, + response: entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_entity_type_with_metadata` + interceptor in new development instead of the `post_get_entity_type` interceptor. + When both interceptors are used, this `post_get_entity_type_with_metadata` interceptor runs after the + `post_get_entity_type` interceptor. The (possibly modified) response returned by + `post_get_entity_type` will be passed to + `post_get_entity_type_with_metadata`. + """ + return response, metadata + def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -561,12 +814,35 @@ def pre_get_feature( def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + def pre_get_featurestore( self, request: featurestore_service.GetFeaturestoreRequest, @@ -587,12 +863,35 @@ def post_get_featurestore( ) -> featurestore.Featurestore: """Post-rpc interceptor for get_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_featurestore` interceptor runs + before the `post_get_featurestore_with_metadata` interceptor. """ return response + def post_get_featurestore_with_metadata( + self, + response: featurestore.Featurestore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[featurestore.Featurestore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_featurestore_with_metadata` + interceptor in new development instead of the `post_get_featurestore` interceptor. + When both interceptors are used, this `post_get_featurestore_with_metadata` interceptor runs after the + `post_get_featurestore` interceptor. The (possibly modified) response returned by + `post_get_featurestore` will be passed to + `post_get_featurestore_with_metadata`. + """ + return response, metadata + def pre_import_feature_values( self, request: featurestore_service.ImportFeatureValuesRequest, @@ -613,12 +912,35 @@ def post_import_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_import_feature_values` interceptor runs + before the `post_import_feature_values_with_metadata` interceptor. """ return response + def post_import_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_import_feature_values_with_metadata` + interceptor in new development instead of the `post_import_feature_values` interceptor. + When both interceptors are used, this `post_import_feature_values_with_metadata` interceptor runs after the + `post_import_feature_values` interceptor. The (possibly modified) response returned by + `post_import_feature_values` will be passed to + `post_import_feature_values_with_metadata`. + """ + return response, metadata + def pre_list_entity_types( self, request: featurestore_service.ListEntityTypesRequest, @@ -639,12 +961,38 @@ def post_list_entity_types( ) -> featurestore_service.ListEntityTypesResponse: """Post-rpc interceptor for list_entity_types - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_entity_types_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_entity_types` interceptor runs + before the `post_list_entity_types_with_metadata` interceptor. """ return response + def post_list_entity_types_with_metadata( + self, + response: featurestore_service.ListEntityTypesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListEntityTypesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_entity_types + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_entity_types_with_metadata` + interceptor in new development instead of the `post_list_entity_types` interceptor. + When both interceptors are used, this `post_list_entity_types_with_metadata` interceptor runs after the + `post_list_entity_types` interceptor. The (possibly modified) response returned by + `post_list_entity_types` will be passed to + `post_list_entity_types_with_metadata`. + """ + return response, metadata + def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -665,12 +1013,38 @@ def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + def pre_list_featurestores( self, request: featurestore_service.ListFeaturestoresRequest, @@ -691,12 +1065,38 @@ def post_list_featurestores( ) -> featurestore_service.ListFeaturestoresResponse: """Post-rpc interceptor for list_featurestores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_featurestores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_featurestores` interceptor runs + before the `post_list_featurestores_with_metadata` interceptor. """ return response + def post_list_featurestores_with_metadata( + self, + response: featurestore_service.ListFeaturestoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturestoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_featurestores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_featurestores_with_metadata` + interceptor in new development instead of the `post_list_featurestores` interceptor. + When both interceptors are used, this `post_list_featurestores_with_metadata` interceptor runs after the + `post_list_featurestores` interceptor. The (possibly modified) response returned by + `post_list_featurestores` will be passed to + `post_list_featurestores_with_metadata`. + """ + return response, metadata + def pre_search_features( self, request: featurestore_service.SearchFeaturesRequest, @@ -717,12 +1117,38 @@ def post_search_features( ) -> featurestore_service.SearchFeaturesResponse: """Post-rpc interceptor for search_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_features` interceptor runs + before the `post_search_features_with_metadata` interceptor. """ return response + def post_search_features_with_metadata( + self, + response: featurestore_service.SearchFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.SearchFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_search_features_with_metadata` + interceptor in new development instead of the `post_search_features` interceptor. + When both interceptors are used, this `post_search_features_with_metadata` interceptor runs after the + `post_search_features` interceptor. The (possibly modified) response returned by + `post_search_features` will be passed to + `post_search_features_with_metadata`. + """ + return response, metadata + def pre_update_entity_type( self, request: featurestore_service.UpdateEntityTypeRequest, @@ -743,12 +1169,35 @@ def post_update_entity_type( ) -> gca_entity_type.EntityType: """Post-rpc interceptor for update_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_entity_type` interceptor runs + before the `post_update_entity_type_with_metadata` interceptor. """ return response + def post_update_entity_type_with_metadata( + self, + response: gca_entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_entity_type_with_metadata` + interceptor in new development instead of the `post_update_entity_type` interceptor. + When both interceptors are used, this `post_update_entity_type_with_metadata` interceptor runs after the + `post_update_entity_type` interceptor. The (possibly modified) response returned by + `post_update_entity_type` will be passed to + `post_update_entity_type_with_metadata`. + """ + return response, metadata + def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -767,12 +1216,35 @@ def pre_update_feature( def post_update_feature(self, response: gca_feature.Feature) -> gca_feature.Feature: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + def post_update_feature_with_metadata( + self, + response: gca_feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + def pre_update_featurestore( self, request: featurestore_service.UpdateFeaturestoreRequest, @@ -793,12 +1265,35 @@ def post_update_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_featurestore` interceptor runs + before the `post_update_featurestore_with_metadata` interceptor. """ return response + def post_update_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_featurestore_with_metadata` + interceptor in new development instead of the `post_update_featurestore` interceptor. + When both interceptors are used, this `post_update_featurestore_with_metadata` interceptor runs after the + `post_update_featurestore` interceptor. The (possibly modified) response returned by + `post_update_featurestore` will be passed to + `post_update_featurestore_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3178,6 +3673,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3332,6 +3831,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3484,6 +3987,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3636,6 +4143,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3790,6 +4301,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3936,6 +4451,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4082,6 +4601,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4230,6 +4753,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4384,6 +4911,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4538,6 +5069,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4687,6 +5222,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4835,6 +5374,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4985,6 +5528,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5139,6 +5686,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5286,6 +5837,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_entity_types(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_entity_types_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5437,6 +5992,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5588,6 +6147,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_featurestores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_featurestores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5735,6 +6298,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5894,6 +6461,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6048,6 +6619,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6202,6 +6777,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest_asyncio.py index ebb83d63bb..253fe4d26b 100644 --- a/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/featurestore_service/transports/rest_asyncio.py @@ -296,12 +296,35 @@ async def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + async def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + async def pre_batch_read_feature_values( self, request: featurestore_service.BatchReadFeatureValuesRequest, @@ -322,12 +345,35 @@ async def post_batch_read_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_feature_values` interceptor runs + before the `post_batch_read_feature_values_with_metadata` interceptor. """ return response + async def post_batch_read_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_read_feature_values_with_metadata` + interceptor in new development instead of the `post_batch_read_feature_values` interceptor. + When both interceptors are used, this `post_batch_read_feature_values_with_metadata` interceptor runs after the + `post_batch_read_feature_values` interceptor. The (possibly modified) response returned by + `post_batch_read_feature_values` will be passed to + `post_batch_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_create_entity_type( self, request: featurestore_service.CreateEntityTypeRequest, @@ -348,12 +394,35 @@ async def post_create_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_entity_type` interceptor runs + before the `post_create_entity_type_with_metadata` interceptor. """ return response + async def post_create_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_entity_type_with_metadata` + interceptor in new development instead of the `post_create_entity_type` interceptor. + When both interceptors are used, this `post_create_entity_type_with_metadata` interceptor runs after the + `post_create_entity_type` interceptor. The (possibly modified) response returned by + `post_create_entity_type` will be passed to + `post_create_entity_type_with_metadata`. + """ + return response, metadata + async def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -374,12 +443,35 @@ async def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + async def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + async def pre_create_featurestore( self, request: featurestore_service.CreateFeaturestoreRequest, @@ -400,12 +492,35 @@ async def post_create_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_featurestore` interceptor runs + before the `post_create_featurestore_with_metadata` interceptor. """ return response + async def post_create_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_featurestore_with_metadata` + interceptor in new development instead of the `post_create_featurestore` interceptor. + When both interceptors are used, this `post_create_featurestore_with_metadata` interceptor runs after the + `post_create_featurestore` interceptor. The (possibly modified) response returned by + `post_create_featurestore` will be passed to + `post_create_featurestore_with_metadata`. + """ + return response, metadata + async def pre_delete_entity_type( self, request: featurestore_service.DeleteEntityTypeRequest, @@ -426,12 +541,35 @@ async def post_delete_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_entity_type` interceptor runs + before the `post_delete_entity_type_with_metadata` interceptor. """ return response + async def post_delete_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_entity_type_with_metadata` + interceptor in new development instead of the `post_delete_entity_type` interceptor. + When both interceptors are used, this `post_delete_entity_type_with_metadata` interceptor runs after the + `post_delete_entity_type` interceptor. The (possibly modified) response returned by + `post_delete_entity_type` will be passed to + `post_delete_entity_type_with_metadata`. + """ + return response, metadata + async def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -452,12 +590,35 @@ async def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + async def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + async def pre_delete_featurestore( self, request: featurestore_service.DeleteFeaturestoreRequest, @@ -478,12 +639,35 @@ async def post_delete_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_featurestore` interceptor runs + before the `post_delete_featurestore_with_metadata` interceptor. """ return response + async def post_delete_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_featurestore_with_metadata` + interceptor in new development instead of the `post_delete_featurestore` interceptor. + When both interceptors are used, this `post_delete_featurestore_with_metadata` interceptor runs after the + `post_delete_featurestore` interceptor. The (possibly modified) response returned by + `post_delete_featurestore` will be passed to + `post_delete_featurestore_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_values( self, request: featurestore_service.DeleteFeatureValuesRequest, @@ -504,12 +688,35 @@ async def post_delete_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_values` interceptor runs + before the `post_delete_feature_values_with_metadata` interceptor. """ return response + async def post_delete_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_values_with_metadata` + interceptor in new development instead of the `post_delete_feature_values` interceptor. + When both interceptors are used, this `post_delete_feature_values_with_metadata` interceptor runs after the + `post_delete_feature_values` interceptor. The (possibly modified) response returned by + `post_delete_feature_values` will be passed to + `post_delete_feature_values_with_metadata`. + """ + return response, metadata + async def pre_export_feature_values( self, request: featurestore_service.ExportFeatureValuesRequest, @@ -530,12 +737,35 @@ async def post_export_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_export_feature_values` interceptor runs + before the `post_export_feature_values_with_metadata` interceptor. """ return response + async def post_export_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_export_feature_values_with_metadata` + interceptor in new development instead of the `post_export_feature_values` interceptor. + When both interceptors are used, this `post_export_feature_values_with_metadata` interceptor runs after the + `post_export_feature_values` interceptor. The (possibly modified) response returned by + `post_export_feature_values` will be passed to + `post_export_feature_values_with_metadata`. + """ + return response, metadata + async def pre_get_entity_type( self, request: featurestore_service.GetEntityTypeRequest, @@ -556,12 +786,35 @@ async def post_get_entity_type( ) -> entity_type.EntityType: """Post-rpc interceptor for get_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_entity_type` interceptor runs + before the `post_get_entity_type_with_metadata` interceptor. """ return response + async def post_get_entity_type_with_metadata( + self, + response: entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_entity_type_with_metadata` + interceptor in new development instead of the `post_get_entity_type` interceptor. + When both interceptors are used, this `post_get_entity_type_with_metadata` interceptor runs after the + `post_get_entity_type` interceptor. The (possibly modified) response returned by + `post_get_entity_type` will be passed to + `post_get_entity_type_with_metadata`. + """ + return response, metadata + async def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -579,12 +832,35 @@ async def pre_get_feature( async def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + async def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + async def pre_get_featurestore( self, request: featurestore_service.GetFeaturestoreRequest, @@ -605,12 +881,35 @@ async def post_get_featurestore( ) -> featurestore.Featurestore: """Post-rpc interceptor for get_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_featurestore` interceptor runs + before the `post_get_featurestore_with_metadata` interceptor. """ return response + async def post_get_featurestore_with_metadata( + self, + response: featurestore.Featurestore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[featurestore.Featurestore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_featurestore_with_metadata` + interceptor in new development instead of the `post_get_featurestore` interceptor. + When both interceptors are used, this `post_get_featurestore_with_metadata` interceptor runs after the + `post_get_featurestore` interceptor. The (possibly modified) response returned by + `post_get_featurestore` will be passed to + `post_get_featurestore_with_metadata`. + """ + return response, metadata + async def pre_import_feature_values( self, request: featurestore_service.ImportFeatureValuesRequest, @@ -631,12 +930,35 @@ async def post_import_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_import_feature_values` interceptor runs + before the `post_import_feature_values_with_metadata` interceptor. """ return response + async def post_import_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_import_feature_values_with_metadata` + interceptor in new development instead of the `post_import_feature_values` interceptor. + When both interceptors are used, this `post_import_feature_values_with_metadata` interceptor runs after the + `post_import_feature_values` interceptor. The (possibly modified) response returned by + `post_import_feature_values` will be passed to + `post_import_feature_values_with_metadata`. + """ + return response, metadata + async def pre_list_entity_types( self, request: featurestore_service.ListEntityTypesRequest, @@ -657,12 +979,38 @@ async def post_list_entity_types( ) -> featurestore_service.ListEntityTypesResponse: """Post-rpc interceptor for list_entity_types - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_entity_types_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_entity_types` interceptor runs + before the `post_list_entity_types_with_metadata` interceptor. """ return response + async def post_list_entity_types_with_metadata( + self, + response: featurestore_service.ListEntityTypesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListEntityTypesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_entity_types + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_entity_types_with_metadata` + interceptor in new development instead of the `post_list_entity_types` interceptor. + When both interceptors are used, this `post_list_entity_types_with_metadata` interceptor runs after the + `post_list_entity_types` interceptor. The (possibly modified) response returned by + `post_list_entity_types` will be passed to + `post_list_entity_types_with_metadata`. + """ + return response, metadata + async def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -683,12 +1031,38 @@ async def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + async def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + async def pre_list_featurestores( self, request: featurestore_service.ListFeaturestoresRequest, @@ -709,12 +1083,38 @@ async def post_list_featurestores( ) -> featurestore_service.ListFeaturestoresResponse: """Post-rpc interceptor for list_featurestores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_featurestores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_featurestores` interceptor runs + before the `post_list_featurestores_with_metadata` interceptor. """ return response + async def post_list_featurestores_with_metadata( + self, + response: featurestore_service.ListFeaturestoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturestoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_featurestores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_featurestores_with_metadata` + interceptor in new development instead of the `post_list_featurestores` interceptor. + When both interceptors are used, this `post_list_featurestores_with_metadata` interceptor runs after the + `post_list_featurestores` interceptor. The (possibly modified) response returned by + `post_list_featurestores` will be passed to + `post_list_featurestores_with_metadata`. + """ + return response, metadata + async def pre_search_features( self, request: featurestore_service.SearchFeaturesRequest, @@ -735,12 +1135,38 @@ async def post_search_features( ) -> featurestore_service.SearchFeaturesResponse: """Post-rpc interceptor for search_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_features` interceptor runs + before the `post_search_features_with_metadata` interceptor. """ return response + async def post_search_features_with_metadata( + self, + response: featurestore_service.SearchFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.SearchFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_search_features_with_metadata` + interceptor in new development instead of the `post_search_features` interceptor. + When both interceptors are used, this `post_search_features_with_metadata` interceptor runs after the + `post_search_features` interceptor. The (possibly modified) response returned by + `post_search_features` will be passed to + `post_search_features_with_metadata`. + """ + return response, metadata + async def pre_update_entity_type( self, request: featurestore_service.UpdateEntityTypeRequest, @@ -761,12 +1187,35 @@ async def post_update_entity_type( ) -> gca_entity_type.EntityType: """Post-rpc interceptor for update_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_entity_type` interceptor runs + before the `post_update_entity_type_with_metadata` interceptor. """ return response + async def post_update_entity_type_with_metadata( + self, + response: gca_entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_entity_type_with_metadata` + interceptor in new development instead of the `post_update_entity_type` interceptor. + When both interceptors are used, this `post_update_entity_type_with_metadata` interceptor runs after the + `post_update_entity_type` interceptor. The (possibly modified) response returned by + `post_update_entity_type` will be passed to + `post_update_entity_type_with_metadata`. + """ + return response, metadata + async def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -787,12 +1236,35 @@ async def post_update_feature( ) -> gca_feature.Feature: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + async def post_update_feature_with_metadata( + self, + response: gca_feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + async def pre_update_featurestore( self, request: featurestore_service.UpdateFeaturestoreRequest, @@ -813,12 +1285,35 @@ async def post_update_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_featurestore` interceptor runs + before the `post_update_featurestore_with_metadata` interceptor. """ return response + async def post_update_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_featurestore_with_metadata` + interceptor in new development instead of the `post_update_featurestore` interceptor. + When both interceptors are used, this `post_update_featurestore_with_metadata` interceptor runs after the + `post_update_featurestore` interceptor. The (possibly modified) response returned by + `post_update_featurestore` will be passed to + `post_update_featurestore_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1435,6 +1930,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1595,6 +2094,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1755,6 +2261,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1917,6 +2427,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2077,6 +2591,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2231,6 +2749,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2387,6 +2909,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2541,6 +3067,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2701,6 +3231,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2861,6 +3395,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3018,6 +3556,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3176,6 +3718,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3332,6 +3878,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3492,6 +4042,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3645,6 +4199,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_entity_types(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_entity_types_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3806,6 +4364,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3961,6 +4523,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_featurestores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_featurestores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4116,6 +4682,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_search_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4281,6 +4851,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4443,6 +5017,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4603,6 +5181,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/client.py b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/client.py index 1ca71f018f..ebc53075be 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/client.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -530,6 +532,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1377,16 +1406,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1432,16 +1465,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1604,16 +1641,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1725,16 +1766,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1847,16 +1892,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1907,16 +1956,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1962,16 +2015,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2017,16 +2074,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest.py b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest.py index bcc8c1eed2..161e9bd274 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest.py @@ -143,12 +143,37 @@ def post_create_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for create_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_create_cached_content` interceptor runs + before the `post_create_cached_content_with_metadata` interceptor. """ return response + def post_create_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_create_cached_content_with_metadata` + interceptor in new development instead of the `post_create_cached_content` interceptor. + When both interceptors are used, this `post_create_cached_content_with_metadata` interceptor runs after the + `post_create_cached_content` interceptor. The (possibly modified) response returned by + `post_create_cached_content` will be passed to + `post_create_cached_content_with_metadata`. + """ + return response, metadata + def pre_delete_cached_content( self, request: gen_ai_cache_service.DeleteCachedContentRequest, @@ -184,12 +209,35 @@ def post_get_cached_content( ) -> cached_content.CachedContent: """Post-rpc interceptor for get_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_get_cached_content` interceptor runs + before the `post_get_cached_content_with_metadata` interceptor. """ return response + def post_get_cached_content_with_metadata( + self, + response: cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_get_cached_content_with_metadata` + interceptor in new development instead of the `post_get_cached_content` interceptor. + When both interceptors are used, this `post_get_cached_content_with_metadata` interceptor runs after the + `post_get_cached_content` interceptor. The (possibly modified) response returned by + `post_get_cached_content` will be passed to + `post_get_cached_content_with_metadata`. + """ + return response, metadata + def pre_list_cached_contents( self, request: gen_ai_cache_service.ListCachedContentsRequest, @@ -210,12 +258,38 @@ def post_list_cached_contents( ) -> gen_ai_cache_service.ListCachedContentsResponse: """Post-rpc interceptor for list_cached_contents - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_cached_contents_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_list_cached_contents` interceptor runs + before the `post_list_cached_contents_with_metadata` interceptor. """ return response + def post_list_cached_contents_with_metadata( + self, + response: gen_ai_cache_service.ListCachedContentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gen_ai_cache_service.ListCachedContentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_cached_contents + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_list_cached_contents_with_metadata` + interceptor in new development instead of the `post_list_cached_contents` interceptor. + When both interceptors are used, this `post_list_cached_contents_with_metadata` interceptor runs after the + `post_list_cached_contents` interceptor. The (possibly modified) response returned by + `post_list_cached_contents` will be passed to + `post_list_cached_contents_with_metadata`. + """ + return response, metadata + def pre_update_cached_content( self, request: gen_ai_cache_service.UpdateCachedContentRequest, @@ -236,12 +310,37 @@ def post_update_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for update_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_update_cached_content` interceptor runs + before the `post_update_cached_content_with_metadata` interceptor. """ return response + def post_update_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_update_cached_content_with_metadata` + interceptor in new development instead of the `post_update_cached_content` interceptor. + When both interceptors are used, this `post_update_cached_content_with_metadata` interceptor runs after the + `post_update_cached_content` interceptor. The (possibly modified) response returned by + `post_update_cached_content` will be passed to + `post_update_cached_content_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -705,6 +804,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -968,6 +1071,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1114,6 +1221,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_cached_contents(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_cached_contents_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1275,6 +1386,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest_asyncio.py index bac94e0043..d8bf862d2e 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_cache_service/transports/rest_asyncio.py @@ -160,12 +160,37 @@ async def post_create_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for create_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_create_cached_content` interceptor runs + before the `post_create_cached_content_with_metadata` interceptor. """ return response + async def post_create_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_create_cached_content_with_metadata` + interceptor in new development instead of the `post_create_cached_content` interceptor. + When both interceptors are used, this `post_create_cached_content_with_metadata` interceptor runs after the + `post_create_cached_content` interceptor. The (possibly modified) response returned by + `post_create_cached_content` will be passed to + `post_create_cached_content_with_metadata`. + """ + return response, metadata + async def pre_delete_cached_content( self, request: gen_ai_cache_service.DeleteCachedContentRequest, @@ -201,12 +226,35 @@ async def post_get_cached_content( ) -> cached_content.CachedContent: """Post-rpc interceptor for get_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_get_cached_content` interceptor runs + before the `post_get_cached_content_with_metadata` interceptor. """ return response + async def post_get_cached_content_with_metadata( + self, + response: cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_get_cached_content_with_metadata` + interceptor in new development instead of the `post_get_cached_content` interceptor. + When both interceptors are used, this `post_get_cached_content_with_metadata` interceptor runs after the + `post_get_cached_content` interceptor. The (possibly modified) response returned by + `post_get_cached_content` will be passed to + `post_get_cached_content_with_metadata`. + """ + return response, metadata + async def pre_list_cached_contents( self, request: gen_ai_cache_service.ListCachedContentsRequest, @@ -227,12 +275,38 @@ async def post_list_cached_contents( ) -> gen_ai_cache_service.ListCachedContentsResponse: """Post-rpc interceptor for list_cached_contents - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_cached_contents_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_list_cached_contents` interceptor runs + before the `post_list_cached_contents_with_metadata` interceptor. """ return response + async def post_list_cached_contents_with_metadata( + self, + response: gen_ai_cache_service.ListCachedContentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gen_ai_cache_service.ListCachedContentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_cached_contents + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_list_cached_contents_with_metadata` + interceptor in new development instead of the `post_list_cached_contents` interceptor. + When both interceptors are used, this `post_list_cached_contents_with_metadata` interceptor runs after the + `post_list_cached_contents` interceptor. The (possibly modified) response returned by + `post_list_cached_contents` will be passed to + `post_list_cached_contents_with_metadata`. + """ + return response, metadata + async def pre_update_cached_content( self, request: gen_ai_cache_service.UpdateCachedContentRequest, @@ -253,12 +327,37 @@ async def post_update_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for update_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_update_cached_content` interceptor runs + before the `post_update_cached_content_with_metadata` interceptor. """ return response + async def post_update_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_update_cached_content_with_metadata` + interceptor in new development instead of the `post_update_cached_content` interceptor. + When both interceptors are used, this `post_update_cached_content_with_metadata` interceptor runs after the + `post_update_cached_content` interceptor. The (possibly modified) response returned by + `post_update_cached_content` will be passed to + `post_update_cached_content_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -789,6 +888,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1062,6 +1165,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1214,6 +1321,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_cached_contents(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_cached_contents_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1379,6 +1490,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/client.py b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/client.py index 358b1687ea..dfbf37b782 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/client.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -599,6 +601,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1477,16 +1506,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1532,16 +1565,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1704,16 +1741,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1825,16 +1866,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1947,16 +1992,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2007,16 +2056,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2062,16 +2115,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2117,16 +2174,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest.py b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest.py index 8fdcdf9061..8cc39be1e9 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest.py @@ -159,12 +159,35 @@ def post_create_tuning_job( ) -> gca_tuning_job.TuningJob: """Post-rpc interceptor for create_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_create_tuning_job` interceptor runs + before the `post_create_tuning_job_with_metadata` interceptor. """ return response + def post_create_tuning_job_with_metadata( + self, + response: gca_tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_create_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_tuning_job` interceptor. + When both interceptors are used, this `post_create_tuning_job_with_metadata` interceptor runs after the + `post_create_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_tuning_job` will be passed to + `post_create_tuning_job_with_metadata`. + """ + return response, metadata + def pre_get_tuning_job( self, request: genai_tuning_service.GetTuningJobRequest, @@ -185,12 +208,35 @@ def post_get_tuning_job( ) -> tuning_job.TuningJob: """Post-rpc interceptor for get_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_get_tuning_job` interceptor runs + before the `post_get_tuning_job_with_metadata` interceptor. """ return response + def post_get_tuning_job_with_metadata( + self, + response: tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_get_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_tuning_job` interceptor. + When both interceptors are used, this `post_get_tuning_job_with_metadata` interceptor runs after the + `post_get_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_tuning_job` will be passed to + `post_get_tuning_job_with_metadata`. + """ + return response, metadata + def pre_list_tuning_jobs( self, request: genai_tuning_service.ListTuningJobsRequest, @@ -211,12 +257,38 @@ def post_list_tuning_jobs( ) -> genai_tuning_service.ListTuningJobsResponse: """Post-rpc interceptor for list_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_list_tuning_jobs` interceptor runs + before the `post_list_tuning_jobs_with_metadata` interceptor. """ return response + def post_list_tuning_jobs_with_metadata( + self, + response: genai_tuning_service.ListTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + genai_tuning_service.ListTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_list_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_tuning_jobs_with_metadata` interceptor runs after the + `post_list_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_tuning_jobs` will be passed to + `post_list_tuning_jobs_with_metadata`. + """ + return response, metadata + def pre_rebase_tuned_model( self, request: genai_tuning_service.RebaseTunedModelRequest, @@ -237,12 +309,35 @@ def post_rebase_tuned_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for rebase_tuned_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_rebase_tuned_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_rebase_tuned_model` interceptor runs + before the `post_rebase_tuned_model_with_metadata` interceptor. """ return response + def post_rebase_tuned_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for rebase_tuned_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_rebase_tuned_model_with_metadata` + interceptor in new development instead of the `post_rebase_tuned_model` interceptor. + When both interceptors are used, this `post_rebase_tuned_model_with_metadata` interceptor runs after the + `post_rebase_tuned_model` interceptor. The (possibly modified) response returned by + `post_rebase_tuned_model` will be passed to + `post_rebase_tuned_model_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2735,6 +2830,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2880,6 +2979,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3027,6 +3130,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3181,6 +3288,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_rebase_tuned_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_rebase_tuned_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest_asyncio.py index b8227da86c..ebc55113ce 100644 --- a/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/gen_ai_tuning_service/transports/rest_asyncio.py @@ -177,12 +177,35 @@ async def post_create_tuning_job( ) -> gca_tuning_job.TuningJob: """Post-rpc interceptor for create_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_create_tuning_job` interceptor runs + before the `post_create_tuning_job_with_metadata` interceptor. """ return response + async def post_create_tuning_job_with_metadata( + self, + response: gca_tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_create_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_tuning_job` interceptor. + When both interceptors are used, this `post_create_tuning_job_with_metadata` interceptor runs after the + `post_create_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_tuning_job` will be passed to + `post_create_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_get_tuning_job( self, request: genai_tuning_service.GetTuningJobRequest, @@ -203,12 +226,35 @@ async def post_get_tuning_job( ) -> tuning_job.TuningJob: """Post-rpc interceptor for get_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_get_tuning_job` interceptor runs + before the `post_get_tuning_job_with_metadata` interceptor. """ return response + async def post_get_tuning_job_with_metadata( + self, + response: tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_get_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_tuning_job` interceptor. + When both interceptors are used, this `post_get_tuning_job_with_metadata` interceptor runs after the + `post_get_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_tuning_job` will be passed to + `post_get_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_list_tuning_jobs( self, request: genai_tuning_service.ListTuningJobsRequest, @@ -229,12 +275,38 @@ async def post_list_tuning_jobs( ) -> genai_tuning_service.ListTuningJobsResponse: """Post-rpc interceptor for list_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_list_tuning_jobs` interceptor runs + before the `post_list_tuning_jobs_with_metadata` interceptor. """ return response + async def post_list_tuning_jobs_with_metadata( + self, + response: genai_tuning_service.ListTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + genai_tuning_service.ListTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_list_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_tuning_jobs_with_metadata` interceptor runs after the + `post_list_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_tuning_jobs` will be passed to + `post_list_tuning_jobs_with_metadata`. + """ + return response, metadata + async def pre_rebase_tuned_model( self, request: genai_tuning_service.RebaseTunedModelRequest, @@ -255,12 +327,35 @@ async def post_rebase_tuned_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for rebase_tuned_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_rebase_tuned_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_rebase_tuned_model` interceptor runs + before the `post_rebase_tuned_model_with_metadata` interceptor. """ return response + async def post_rebase_tuned_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for rebase_tuned_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_rebase_tuned_model_with_metadata` + interceptor in new development instead of the `post_rebase_tuned_model` interceptor. + When both interceptors are used, this `post_rebase_tuned_model_with_metadata` interceptor runs after the + `post_rebase_tuned_model` interceptor. The (possibly modified) response returned by + `post_rebase_tuned_model` will be passed to + `post_rebase_tuned_model_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -916,6 +1011,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1071,6 +1170,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1224,6 +1327,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1386,6 +1493,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_rebase_tuned_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_rebase_tuned_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/index_endpoint_service/client.py b/google/cloud/aiplatform_v1/services/index_endpoint_service/client.py index c392d30ce9..15ddef8c84 100644 --- a/google/cloud/aiplatform_v1/services/index_endpoint_service/client.py +++ b/google/cloud/aiplatform_v1/services/index_endpoint_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -555,6 +557,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1880,16 +1909,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1935,16 +1968,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2107,16 +2144,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2228,16 +2269,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2350,16 +2395,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2410,16 +2459,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2465,16 +2518,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2520,16 +2577,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest.py b/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest.py index f3aeb36310..34896fa38e 100644 --- a/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest.py @@ -171,12 +171,35 @@ def post_create_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_index_endpoint` interceptor runs + before the `post_create_index_endpoint_with_metadata` interceptor. """ return response + def post_create_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_create_index_endpoint_with_metadata` + interceptor in new development instead of the `post_create_index_endpoint` interceptor. + When both interceptors are used, this `post_create_index_endpoint_with_metadata` interceptor runs after the + `post_create_index_endpoint` interceptor. The (possibly modified) response returned by + `post_create_index_endpoint` will be passed to + `post_create_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_delete_index_endpoint( self, request: index_endpoint_service.DeleteIndexEndpointRequest, @@ -197,12 +220,35 @@ def post_delete_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index_endpoint` interceptor runs + before the `post_delete_index_endpoint_with_metadata` interceptor. """ return response + def post_delete_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_index_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_index_endpoint` interceptor. + When both interceptors are used, this `post_delete_index_endpoint_with_metadata` interceptor runs after the + `post_delete_index_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_index_endpoint` will be passed to + `post_delete_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_deploy_index( self, request: index_endpoint_service.DeployIndexRequest, @@ -223,12 +269,35 @@ def post_deploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_index` interceptor runs + before the `post_deploy_index_with_metadata` interceptor. """ return response + def post_deploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_index_with_metadata` + interceptor in new development instead of the `post_deploy_index` interceptor. + When both interceptors are used, this `post_deploy_index_with_metadata` interceptor runs after the + `post_deploy_index` interceptor. The (possibly modified) response returned by + `post_deploy_index` will be passed to + `post_deploy_index_with_metadata`. + """ + return response, metadata + def pre_get_index_endpoint( self, request: index_endpoint_service.GetIndexEndpointRequest, @@ -249,12 +318,35 @@ def post_get_index_endpoint( ) -> index_endpoint.IndexEndpoint: """Post-rpc interceptor for get_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_index_endpoint` interceptor runs + before the `post_get_index_endpoint_with_metadata` interceptor. """ return response + def post_get_index_endpoint_with_metadata( + self, + response: index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_get_index_endpoint_with_metadata` + interceptor in new development instead of the `post_get_index_endpoint` interceptor. + When both interceptors are used, this `post_get_index_endpoint_with_metadata` interceptor runs after the + `post_get_index_endpoint` interceptor. The (possibly modified) response returned by + `post_get_index_endpoint` will be passed to + `post_get_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_list_index_endpoints( self, request: index_endpoint_service.ListIndexEndpointsRequest, @@ -275,12 +367,38 @@ def post_list_index_endpoints( ) -> index_endpoint_service.ListIndexEndpointsResponse: """Post-rpc interceptor for list_index_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_index_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_index_endpoints` interceptor runs + before the `post_list_index_endpoints_with_metadata` interceptor. """ return response + def post_list_index_endpoints_with_metadata( + self, + response: index_endpoint_service.ListIndexEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_endpoint_service.ListIndexEndpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_index_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_list_index_endpoints_with_metadata` + interceptor in new development instead of the `post_list_index_endpoints` interceptor. + When both interceptors are used, this `post_list_index_endpoints_with_metadata` interceptor runs after the + `post_list_index_endpoints` interceptor. The (possibly modified) response returned by + `post_list_index_endpoints` will be passed to + `post_list_index_endpoints_with_metadata`. + """ + return response, metadata + def pre_mutate_deployed_index( self, request: index_endpoint_service.MutateDeployedIndexRequest, @@ -301,12 +419,35 @@ def post_mutate_deployed_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_index` interceptor runs + before the `post_mutate_deployed_index_with_metadata` interceptor. """ return response + def post_mutate_deployed_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_index_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_index` interceptor. + When both interceptors are used, this `post_mutate_deployed_index_with_metadata` interceptor runs after the + `post_mutate_deployed_index` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_index` will be passed to + `post_mutate_deployed_index_with_metadata`. + """ + return response, metadata + def pre_undeploy_index( self, request: index_endpoint_service.UndeployIndexRequest, @@ -327,12 +468,35 @@ def post_undeploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_index` interceptor runs + before the `post_undeploy_index_with_metadata` interceptor. """ return response + def post_undeploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_index_with_metadata` + interceptor in new development instead of the `post_undeploy_index` interceptor. + When both interceptors are used, this `post_undeploy_index_with_metadata` interceptor runs after the + `post_undeploy_index` interceptor. The (possibly modified) response returned by + `post_undeploy_index` will be passed to + `post_undeploy_index_with_metadata`. + """ + return response, metadata + def pre_update_index_endpoint( self, request: index_endpoint_service.UpdateIndexEndpointRequest, @@ -353,12 +517,37 @@ def post_update_index_endpoint( ) -> gca_index_endpoint.IndexEndpoint: """Post-rpc interceptor for update_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_index_endpoint` interceptor runs + before the `post_update_index_endpoint_with_metadata` interceptor. """ return response + def post_update_index_endpoint_with_metadata( + self, + response: gca_index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_update_index_endpoint_with_metadata` + interceptor in new development instead of the `post_update_index_endpoint` interceptor. + When both interceptors are used, this `post_update_index_endpoint_with_metadata` interceptor runs after the + `post_update_index_endpoint` interceptor. The (possibly modified) response returned by + `post_update_index_endpoint` will be passed to + `post_update_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2735,6 +2924,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2883,6 +3076,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3033,6 +3230,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_deploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_deploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3183,6 +3384,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3332,6 +3537,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_index_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_index_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3490,6 +3699,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_mutate_deployed_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_mutate_deployed_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3640,6 +3853,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_undeploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_undeploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3796,6 +4013,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest_asyncio.py index f2a4c14fbf..ad455e7abb 100644 --- a/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/index_endpoint_service/transports/rest_asyncio.py @@ -189,12 +189,35 @@ async def post_create_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_index_endpoint` interceptor runs + before the `post_create_index_endpoint_with_metadata` interceptor. """ return response + async def post_create_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_create_index_endpoint_with_metadata` + interceptor in new development instead of the `post_create_index_endpoint` interceptor. + When both interceptors are used, this `post_create_index_endpoint_with_metadata` interceptor runs after the + `post_create_index_endpoint` interceptor. The (possibly modified) response returned by + `post_create_index_endpoint` will be passed to + `post_create_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_delete_index_endpoint( self, request: index_endpoint_service.DeleteIndexEndpointRequest, @@ -215,12 +238,35 @@ async def post_delete_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index_endpoint` interceptor runs + before the `post_delete_index_endpoint_with_metadata` interceptor. """ return response + async def post_delete_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_index_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_index_endpoint` interceptor. + When both interceptors are used, this `post_delete_index_endpoint_with_metadata` interceptor runs after the + `post_delete_index_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_index_endpoint` will be passed to + `post_delete_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_deploy_index( self, request: index_endpoint_service.DeployIndexRequest, @@ -241,12 +287,35 @@ async def post_deploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_index` interceptor runs + before the `post_deploy_index_with_metadata` interceptor. """ return response + async def post_deploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_index_with_metadata` + interceptor in new development instead of the `post_deploy_index` interceptor. + When both interceptors are used, this `post_deploy_index_with_metadata` interceptor runs after the + `post_deploy_index` interceptor. The (possibly modified) response returned by + `post_deploy_index` will be passed to + `post_deploy_index_with_metadata`. + """ + return response, metadata + async def pre_get_index_endpoint( self, request: index_endpoint_service.GetIndexEndpointRequest, @@ -267,12 +336,35 @@ async def post_get_index_endpoint( ) -> index_endpoint.IndexEndpoint: """Post-rpc interceptor for get_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_index_endpoint` interceptor runs + before the `post_get_index_endpoint_with_metadata` interceptor. """ return response + async def post_get_index_endpoint_with_metadata( + self, + response: index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_get_index_endpoint_with_metadata` + interceptor in new development instead of the `post_get_index_endpoint` interceptor. + When both interceptors are used, this `post_get_index_endpoint_with_metadata` interceptor runs after the + `post_get_index_endpoint` interceptor. The (possibly modified) response returned by + `post_get_index_endpoint` will be passed to + `post_get_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_list_index_endpoints( self, request: index_endpoint_service.ListIndexEndpointsRequest, @@ -293,12 +385,38 @@ async def post_list_index_endpoints( ) -> index_endpoint_service.ListIndexEndpointsResponse: """Post-rpc interceptor for list_index_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_index_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_index_endpoints` interceptor runs + before the `post_list_index_endpoints_with_metadata` interceptor. """ return response + async def post_list_index_endpoints_with_metadata( + self, + response: index_endpoint_service.ListIndexEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_endpoint_service.ListIndexEndpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_index_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_list_index_endpoints_with_metadata` + interceptor in new development instead of the `post_list_index_endpoints` interceptor. + When both interceptors are used, this `post_list_index_endpoints_with_metadata` interceptor runs after the + `post_list_index_endpoints` interceptor. The (possibly modified) response returned by + `post_list_index_endpoints` will be passed to + `post_list_index_endpoints_with_metadata`. + """ + return response, metadata + async def pre_mutate_deployed_index( self, request: index_endpoint_service.MutateDeployedIndexRequest, @@ -319,12 +437,35 @@ async def post_mutate_deployed_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_index` interceptor runs + before the `post_mutate_deployed_index_with_metadata` interceptor. """ return response + async def post_mutate_deployed_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_index_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_index` interceptor. + When both interceptors are used, this `post_mutate_deployed_index_with_metadata` interceptor runs after the + `post_mutate_deployed_index` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_index` will be passed to + `post_mutate_deployed_index_with_metadata`. + """ + return response, metadata + async def pre_undeploy_index( self, request: index_endpoint_service.UndeployIndexRequest, @@ -345,12 +486,35 @@ async def post_undeploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_index` interceptor runs + before the `post_undeploy_index_with_metadata` interceptor. """ return response + async def post_undeploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_index_with_metadata` + interceptor in new development instead of the `post_undeploy_index` interceptor. + When both interceptors are used, this `post_undeploy_index_with_metadata` interceptor runs after the + `post_undeploy_index` interceptor. The (possibly modified) response returned by + `post_undeploy_index` will be passed to + `post_undeploy_index_with_metadata`. + """ + return response, metadata + async def pre_update_index_endpoint( self, request: index_endpoint_service.UpdateIndexEndpointRequest, @@ -371,12 +535,37 @@ async def post_update_index_endpoint( ) -> gca_index_endpoint.IndexEndpoint: """Post-rpc interceptor for update_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_index_endpoint` interceptor runs + before the `post_update_index_endpoint_with_metadata` interceptor. """ return response + async def post_update_index_endpoint_with_metadata( + self, + response: gca_index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_update_index_endpoint_with_metadata` + interceptor in new development instead of the `post_update_index_endpoint` interceptor. + When both interceptors are used, this `post_update_index_endpoint_with_metadata` interceptor runs after the + `post_update_index_endpoint` interceptor. The (possibly modified) response returned by + `post_update_index_endpoint` will be passed to + `post_update_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -925,6 +1114,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1079,6 +1272,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1241,6 +1438,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_deploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_deploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1395,6 +1596,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1548,6 +1753,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_index_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_index_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1712,6 +1921,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_mutate_deployed_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_mutate_deployed_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1872,6 +2085,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_undeploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_undeploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2032,6 +2249,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/index_service/client.py b/google/cloud/aiplatform_v1/services/index_service/client.py index f05d36dbcb..d66a91d446 100644 --- a/google/cloud/aiplatform_v1/services/index_service/client.py +++ b/google/cloud/aiplatform_v1/services/index_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -534,6 +536,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1604,16 +1633,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1659,16 +1692,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1831,16 +1868,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1952,16 +1993,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2074,16 +2119,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2134,16 +2183,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2189,16 +2242,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2244,16 +2301,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/index_service/transports/rest.py b/google/cloud/aiplatform_v1/services/index_service/transports/rest.py index a0ea86eb7a..f9de906acd 100644 --- a/google/cloud/aiplatform_v1/services/index_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/index_service/transports/rest.py @@ -161,12 +161,35 @@ def post_create_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_create_index` interceptor runs + before the `post_create_index_with_metadata` interceptor. """ return response + def post_create_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_create_index_with_metadata` + interceptor in new development instead of the `post_create_index` interceptor. + When both interceptors are used, this `post_create_index_with_metadata` interceptor runs after the + `post_create_index` interceptor. The (possibly modified) response returned by + `post_create_index` will be passed to + `post_create_index_with_metadata`. + """ + return response, metadata + def pre_delete_index( self, request: index_service.DeleteIndexRequest, @@ -186,12 +209,35 @@ def post_delete_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index` interceptor runs + before the `post_delete_index_with_metadata` interceptor. """ return response + def post_delete_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_delete_index_with_metadata` + interceptor in new development instead of the `post_delete_index` interceptor. + When both interceptors are used, this `post_delete_index_with_metadata` interceptor runs after the + `post_delete_index` interceptor. The (possibly modified) response returned by + `post_delete_index` will be passed to + `post_delete_index_with_metadata`. + """ + return response, metadata + def pre_get_index( self, request: index_service.GetIndexRequest, @@ -207,12 +253,33 @@ def pre_get_index( def post_get_index(self, response: index.Index) -> index.Index: """Post-rpc interceptor for get_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_get_index` interceptor runs + before the `post_get_index_with_metadata` interceptor. """ return response + def post_get_index_with_metadata( + self, response: index.Index, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[index.Index, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_get_index_with_metadata` + interceptor in new development instead of the `post_get_index` interceptor. + When both interceptors are used, this `post_get_index_with_metadata` interceptor runs after the + `post_get_index` interceptor. The (possibly modified) response returned by + `post_get_index` will be passed to + `post_get_index_with_metadata`. + """ + return response, metadata + def pre_list_indexes( self, request: index_service.ListIndexesRequest, @@ -232,12 +299,37 @@ def post_list_indexes( ) -> index_service.ListIndexesResponse: """Post-rpc interceptor for list_indexes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_indexes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_list_indexes` interceptor runs + before the `post_list_indexes_with_metadata` interceptor. """ return response + def post_list_indexes_with_metadata( + self, + response: index_service.ListIndexesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.ListIndexesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_indexes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_list_indexes_with_metadata` + interceptor in new development instead of the `post_list_indexes` interceptor. + When both interceptors are used, this `post_list_indexes_with_metadata` interceptor runs after the + `post_list_indexes` interceptor. The (possibly modified) response returned by + `post_list_indexes` will be passed to + `post_list_indexes_with_metadata`. + """ + return response, metadata + def pre_remove_datapoints( self, request: index_service.RemoveDatapointsRequest, @@ -257,12 +349,37 @@ def post_remove_datapoints( ) -> index_service.RemoveDatapointsResponse: """Post-rpc interceptor for remove_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_remove_datapoints` interceptor runs + before the `post_remove_datapoints_with_metadata` interceptor. """ return response + def post_remove_datapoints_with_metadata( + self, + response: index_service.RemoveDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.RemoveDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for remove_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_remove_datapoints_with_metadata` + interceptor in new development instead of the `post_remove_datapoints` interceptor. + When both interceptors are used, this `post_remove_datapoints_with_metadata` interceptor runs after the + `post_remove_datapoints` interceptor. The (possibly modified) response returned by + `post_remove_datapoints` will be passed to + `post_remove_datapoints_with_metadata`. + """ + return response, metadata + def pre_update_index( self, request: index_service.UpdateIndexRequest, @@ -282,12 +399,35 @@ def post_update_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_update_index` interceptor runs + before the `post_update_index_with_metadata` interceptor. """ return response + def post_update_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_update_index_with_metadata` + interceptor in new development instead of the `post_update_index` interceptor. + When both interceptors are used, this `post_update_index_with_metadata` interceptor runs after the + `post_update_index` interceptor. The (possibly modified) response returned by + `post_update_index` will be passed to + `post_update_index_with_metadata`. + """ + return response, metadata + def pre_upsert_datapoints( self, request: index_service.UpsertDatapointsRequest, @@ -307,12 +447,37 @@ def post_upsert_datapoints( ) -> index_service.UpsertDatapointsResponse: """Post-rpc interceptor for upsert_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upsert_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_upsert_datapoints` interceptor runs + before the `post_upsert_datapoints_with_metadata` interceptor. """ return response + def post_upsert_datapoints_with_metadata( + self, + response: index_service.UpsertDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.UpsertDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for upsert_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_upsert_datapoints_with_metadata` + interceptor in new development instead of the `post_upsert_datapoints` interceptor. + When both interceptors are used, this `post_upsert_datapoints_with_metadata` interceptor runs after the + `post_upsert_datapoints` interceptor. The (possibly modified) response returned by + `post_upsert_datapoints` will be passed to + `post_upsert_datapoints_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2691,6 +2856,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2838,6 +3007,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2986,6 +3159,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3134,6 +3311,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_indexes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_indexes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3288,6 +3469,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_remove_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_remove_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3445,6 +3630,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3597,6 +3786,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_upsert_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upsert_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/index_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/index_service/transports/rest_asyncio.py index c22e9bdfae..413707acd1 100644 --- a/google/cloud/aiplatform_v1/services/index_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/index_service/transports/rest_asyncio.py @@ -179,12 +179,35 @@ async def post_create_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_create_index` interceptor runs + before the `post_create_index_with_metadata` interceptor. """ return response + async def post_create_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_create_index_with_metadata` + interceptor in new development instead of the `post_create_index` interceptor. + When both interceptors are used, this `post_create_index_with_metadata` interceptor runs after the + `post_create_index` interceptor. The (possibly modified) response returned by + `post_create_index` will be passed to + `post_create_index_with_metadata`. + """ + return response, metadata + async def pre_delete_index( self, request: index_service.DeleteIndexRequest, @@ -204,12 +227,35 @@ async def post_delete_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index` interceptor runs + before the `post_delete_index_with_metadata` interceptor. """ return response + async def post_delete_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_delete_index_with_metadata` + interceptor in new development instead of the `post_delete_index` interceptor. + When both interceptors are used, this `post_delete_index_with_metadata` interceptor runs after the + `post_delete_index` interceptor. The (possibly modified) response returned by + `post_delete_index` will be passed to + `post_delete_index_with_metadata`. + """ + return response, metadata + async def pre_get_index( self, request: index_service.GetIndexRequest, @@ -225,12 +271,33 @@ async def pre_get_index( async def post_get_index(self, response: index.Index) -> index.Index: """Post-rpc interceptor for get_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_get_index` interceptor runs + before the `post_get_index_with_metadata` interceptor. """ return response + async def post_get_index_with_metadata( + self, response: index.Index, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[index.Index, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_get_index_with_metadata` + interceptor in new development instead of the `post_get_index` interceptor. + When both interceptors are used, this `post_get_index_with_metadata` interceptor runs after the + `post_get_index` interceptor. The (possibly modified) response returned by + `post_get_index` will be passed to + `post_get_index_with_metadata`. + """ + return response, metadata + async def pre_list_indexes( self, request: index_service.ListIndexesRequest, @@ -250,12 +317,37 @@ async def post_list_indexes( ) -> index_service.ListIndexesResponse: """Post-rpc interceptor for list_indexes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_indexes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_list_indexes` interceptor runs + before the `post_list_indexes_with_metadata` interceptor. """ return response + async def post_list_indexes_with_metadata( + self, + response: index_service.ListIndexesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.ListIndexesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_indexes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_list_indexes_with_metadata` + interceptor in new development instead of the `post_list_indexes` interceptor. + When both interceptors are used, this `post_list_indexes_with_metadata` interceptor runs after the + `post_list_indexes` interceptor. The (possibly modified) response returned by + `post_list_indexes` will be passed to + `post_list_indexes_with_metadata`. + """ + return response, metadata + async def pre_remove_datapoints( self, request: index_service.RemoveDatapointsRequest, @@ -275,12 +367,37 @@ async def post_remove_datapoints( ) -> index_service.RemoveDatapointsResponse: """Post-rpc interceptor for remove_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_remove_datapoints` interceptor runs + before the `post_remove_datapoints_with_metadata` interceptor. """ return response + async def post_remove_datapoints_with_metadata( + self, + response: index_service.RemoveDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.RemoveDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for remove_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_remove_datapoints_with_metadata` + interceptor in new development instead of the `post_remove_datapoints` interceptor. + When both interceptors are used, this `post_remove_datapoints_with_metadata` interceptor runs after the + `post_remove_datapoints` interceptor. The (possibly modified) response returned by + `post_remove_datapoints` will be passed to + `post_remove_datapoints_with_metadata`. + """ + return response, metadata + async def pre_update_index( self, request: index_service.UpdateIndexRequest, @@ -300,12 +417,35 @@ async def post_update_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_update_index` interceptor runs + before the `post_update_index_with_metadata` interceptor. """ return response + async def post_update_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_update_index_with_metadata` + interceptor in new development instead of the `post_update_index` interceptor. + When both interceptors are used, this `post_update_index_with_metadata` interceptor runs after the + `post_update_index` interceptor. The (possibly modified) response returned by + `post_update_index` will be passed to + `post_update_index_with_metadata`. + """ + return response, metadata + async def pre_upsert_datapoints( self, request: index_service.UpsertDatapointsRequest, @@ -325,12 +465,37 @@ async def post_upsert_datapoints( ) -> index_service.UpsertDatapointsResponse: """Post-rpc interceptor for upsert_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upsert_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_upsert_datapoints` interceptor runs + before the `post_upsert_datapoints_with_metadata` interceptor. """ return response + async def post_upsert_datapoints_with_metadata( + self, + response: index_service.UpsertDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.UpsertDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for upsert_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_upsert_datapoints_with_metadata` + interceptor in new development instead of the `post_upsert_datapoints` interceptor. + When both interceptors are used, this `post_upsert_datapoints_with_metadata` interceptor runs after the + `post_upsert_datapoints` interceptor. The (possibly modified) response returned by + `post_upsert_datapoints` will be passed to + `post_upsert_datapoints_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -880,6 +1045,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1037,6 +1206,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1193,6 +1366,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1349,6 +1526,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_indexes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_indexes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1511,6 +1692,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_remove_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_remove_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1678,6 +1863,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1838,6 +2027,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upsert_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upsert_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/job_service/client.py b/google/cloud/aiplatform_v1/services/job_service/client.py index 92c13b8bd2..b16f167f0b 100644 --- a/google/cloud/aiplatform_v1/services/job_service/client.py +++ b/google/cloud/aiplatform_v1/services/job_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -892,6 +894,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -5413,16 +5442,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -5468,16 +5501,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -5640,16 +5677,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5761,16 +5802,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5883,16 +5928,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5943,16 +5992,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -5998,16 +6051,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -6053,16 +6110,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/job_service/transports/rest.py b/google/cloud/aiplatform_v1/services/job_service/transports/rest.py index d68005119e..53b106bffa 100644 --- a/google/cloud/aiplatform_v1/services/job_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/job_service/transports/rest.py @@ -449,12 +449,38 @@ def post_create_batch_prediction_job( ) -> gca_batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for create_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_batch_prediction_job` interceptor runs + before the `post_create_batch_prediction_job_with_metadata` interceptor. """ return response + def post_create_batch_prediction_job_with_metadata( + self, + response: gca_batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_batch_prediction_job.BatchPredictionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_create_batch_prediction_job` interceptor. + When both interceptors are used, this `post_create_batch_prediction_job_with_metadata` interceptor runs after the + `post_create_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_create_batch_prediction_job` will be passed to + `post_create_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_create_custom_job( self, request: job_service.CreateCustomJobRequest, @@ -474,12 +500,35 @@ def post_create_custom_job( ) -> gca_custom_job.CustomJob: """Post-rpc interceptor for create_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_custom_job` interceptor runs + before the `post_create_custom_job_with_metadata` interceptor. """ return response + def post_create_custom_job_with_metadata( + self, + response: gca_custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_custom_job_with_metadata` + interceptor in new development instead of the `post_create_custom_job` interceptor. + When both interceptors are used, this `post_create_custom_job_with_metadata` interceptor runs after the + `post_create_custom_job` interceptor. The (possibly modified) response returned by + `post_create_custom_job` will be passed to + `post_create_custom_job_with_metadata`. + """ + return response, metadata + def pre_create_data_labeling_job( self, request: job_service.CreateDataLabelingJobRequest, @@ -500,12 +549,37 @@ def post_create_data_labeling_job( ) -> gca_data_labeling_job.DataLabelingJob: """Post-rpc interceptor for create_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_data_labeling_job` interceptor runs + before the `post_create_data_labeling_job_with_metadata` interceptor. """ return response + def post_create_data_labeling_job_with_metadata( + self, + response: gca_data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_create_data_labeling_job` interceptor. + When both interceptors are used, this `post_create_data_labeling_job_with_metadata` interceptor runs after the + `post_create_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_create_data_labeling_job` will be passed to + `post_create_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_create_hyperparameter_tuning_job( self, request: job_service.CreateHyperparameterTuningJobRequest, @@ -526,12 +600,38 @@ def post_create_hyperparameter_tuning_job( ) -> gca_hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for create_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_hyperparameter_tuning_job` interceptor runs + before the `post_create_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_create_hyperparameter_tuning_job_with_metadata( + self, + response: gca_hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_create_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_create_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_hyperparameter_tuning_job` will be passed to + `post_create_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_create_model_deployment_monitoring_job( self, request: job_service.CreateModelDeploymentMonitoringJobRequest, @@ -552,12 +652,38 @@ def post_create_model_deployment_monitoring_job( ) -> gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for create_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_deployment_monitoring_job` interceptor runs + before the `post_create_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_create_model_deployment_monitoring_job_with_metadata( + self, + response: gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_deployment_monitoring_job` will be passed to + `post_create_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_create_nas_job( self, request: job_service.CreateNasJobRequest, @@ -575,12 +701,35 @@ def pre_create_nas_job( def post_create_nas_job(self, response: gca_nas_job.NasJob) -> gca_nas_job.NasJob: """Post-rpc interceptor for create_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_nas_job` interceptor runs + before the `post_create_nas_job_with_metadata` interceptor. """ return response + def post_create_nas_job_with_metadata( + self, + response: gca_nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_nas_job_with_metadata` + interceptor in new development instead of the `post_create_nas_job` interceptor. + When both interceptors are used, this `post_create_nas_job_with_metadata` interceptor runs after the + `post_create_nas_job` interceptor. The (possibly modified) response returned by + `post_create_nas_job` will be passed to + `post_create_nas_job_with_metadata`. + """ + return response, metadata + def pre_delete_batch_prediction_job( self, request: job_service.DeleteBatchPredictionJobRequest, @@ -601,12 +750,35 @@ def post_delete_batch_prediction_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_batch_prediction_job` interceptor runs + before the `post_delete_batch_prediction_job_with_metadata` interceptor. """ return response + def post_delete_batch_prediction_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_delete_batch_prediction_job` interceptor. + When both interceptors are used, this `post_delete_batch_prediction_job_with_metadata` interceptor runs after the + `post_delete_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_delete_batch_prediction_job` will be passed to + `post_delete_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_delete_custom_job( self, request: job_service.DeleteCustomJobRequest, @@ -626,12 +798,35 @@ def post_delete_custom_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_custom_job` interceptor runs + before the `post_delete_custom_job_with_metadata` interceptor. """ return response + def post_delete_custom_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_custom_job_with_metadata` + interceptor in new development instead of the `post_delete_custom_job` interceptor. + When both interceptors are used, this `post_delete_custom_job_with_metadata` interceptor runs after the + `post_delete_custom_job` interceptor. The (possibly modified) response returned by + `post_delete_custom_job` will be passed to + `post_delete_custom_job_with_metadata`. + """ + return response, metadata + def pre_delete_data_labeling_job( self, request: job_service.DeleteDataLabelingJobRequest, @@ -652,12 +847,35 @@ def post_delete_data_labeling_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_data_labeling_job` interceptor runs + before the `post_delete_data_labeling_job_with_metadata` interceptor. """ return response + def post_delete_data_labeling_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_delete_data_labeling_job` interceptor. + When both interceptors are used, this `post_delete_data_labeling_job_with_metadata` interceptor runs after the + `post_delete_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_delete_data_labeling_job` will be passed to + `post_delete_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_delete_hyperparameter_tuning_job( self, request: job_service.DeleteHyperparameterTuningJobRequest, @@ -678,12 +896,35 @@ def post_delete_hyperparameter_tuning_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_hyperparameter_tuning_job` interceptor runs + before the `post_delete_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_delete_hyperparameter_tuning_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_delete_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_delete_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_delete_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_delete_hyperparameter_tuning_job` will be passed to + `post_delete_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_delete_model_deployment_monitoring_job( self, request: job_service.DeleteModelDeploymentMonitoringJobRequest, @@ -704,12 +945,35 @@ def post_delete_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_deployment_monitoring_job` interceptor runs + before the `post_delete_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_delete_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_deployment_monitoring_job` will be passed to + `post_delete_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_delete_nas_job( self, request: job_service.DeleteNasJobRequest, @@ -729,12 +993,35 @@ def post_delete_nas_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_nas_job` interceptor runs + before the `post_delete_nas_job_with_metadata` interceptor. """ return response + def post_delete_nas_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_nas_job_with_metadata` + interceptor in new development instead of the `post_delete_nas_job` interceptor. + When both interceptors are used, this `post_delete_nas_job_with_metadata` interceptor runs after the + `post_delete_nas_job` interceptor. The (possibly modified) response returned by + `post_delete_nas_job` will be passed to + `post_delete_nas_job_with_metadata`. + """ + return response, metadata + def pre_get_batch_prediction_job( self, request: job_service.GetBatchPredictionJobRequest, @@ -755,12 +1042,37 @@ def post_get_batch_prediction_job( ) -> batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for get_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_batch_prediction_job` interceptor runs + before the `post_get_batch_prediction_job_with_metadata` interceptor. """ return response + def post_get_batch_prediction_job_with_metadata( + self, + response: batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + batch_prediction_job.BatchPredictionJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_get_batch_prediction_job` interceptor. + When both interceptors are used, this `post_get_batch_prediction_job_with_metadata` interceptor runs after the + `post_get_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_get_batch_prediction_job` will be passed to + `post_get_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_get_custom_job( self, request: job_service.GetCustomJobRequest, @@ -780,12 +1092,35 @@ def post_get_custom_job( ) -> custom_job.CustomJob: """Post-rpc interceptor for get_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_custom_job` interceptor runs + before the `post_get_custom_job_with_metadata` interceptor. """ return response + def post_get_custom_job_with_metadata( + self, + response: custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_custom_job_with_metadata` + interceptor in new development instead of the `post_get_custom_job` interceptor. + When both interceptors are used, this `post_get_custom_job_with_metadata` interceptor runs after the + `post_get_custom_job` interceptor. The (possibly modified) response returned by + `post_get_custom_job` will be passed to + `post_get_custom_job_with_metadata`. + """ + return response, metadata + def pre_get_data_labeling_job( self, request: job_service.GetDataLabelingJobRequest, @@ -805,12 +1140,37 @@ def post_get_data_labeling_job( ) -> data_labeling_job.DataLabelingJob: """Post-rpc interceptor for get_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_data_labeling_job` interceptor runs + before the `post_get_data_labeling_job_with_metadata` interceptor. """ return response + def post_get_data_labeling_job_with_metadata( + self, + response: data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_get_data_labeling_job` interceptor. + When both interceptors are used, this `post_get_data_labeling_job_with_metadata` interceptor runs after the + `post_get_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_get_data_labeling_job` will be passed to + `post_get_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_get_hyperparameter_tuning_job( self, request: job_service.GetHyperparameterTuningJobRequest, @@ -831,12 +1191,38 @@ def post_get_hyperparameter_tuning_job( ) -> hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for get_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_hyperparameter_tuning_job` interceptor runs + before the `post_get_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_get_hyperparameter_tuning_job_with_metadata( + self, + response: hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_get_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_get_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_hyperparameter_tuning_job` will be passed to + `post_get_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_get_model_deployment_monitoring_job( self, request: job_service.GetModelDeploymentMonitoringJobRequest, @@ -857,12 +1243,38 @@ def post_get_model_deployment_monitoring_job( ) -> model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for get_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_deployment_monitoring_job` interceptor runs + before the `post_get_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_get_model_deployment_monitoring_job_with_metadata( + self, + response: model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_deployment_monitoring_job` will be passed to + `post_get_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_get_nas_job( self, request: job_service.GetNasJobRequest, @@ -878,12 +1290,35 @@ def pre_get_nas_job( def post_get_nas_job(self, response: nas_job.NasJob) -> nas_job.NasJob: """Post-rpc interceptor for get_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_job` interceptor runs + before the `post_get_nas_job_with_metadata` interceptor. """ return response + def post_get_nas_job_with_metadata( + self, + response: nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_job_with_metadata` + interceptor in new development instead of the `post_get_nas_job` interceptor. + When both interceptors are used, this `post_get_nas_job_with_metadata` interceptor runs after the + `post_get_nas_job` interceptor. The (possibly modified) response returned by + `post_get_nas_job` will be passed to + `post_get_nas_job_with_metadata`. + """ + return response, metadata + def pre_get_nas_trial_detail( self, request: job_service.GetNasTrialDetailRequest, @@ -903,12 +1338,35 @@ def post_get_nas_trial_detail( ) -> nas_job.NasTrialDetail: """Post-rpc interceptor for get_nas_trial_detail - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_trial_detail_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_trial_detail` interceptor runs + before the `post_get_nas_trial_detail_with_metadata` interceptor. """ return response + def post_get_nas_trial_detail_with_metadata( + self, + response: nas_job.NasTrialDetail, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasTrialDetail, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_trial_detail + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_trial_detail_with_metadata` + interceptor in new development instead of the `post_get_nas_trial_detail` interceptor. + When both interceptors are used, this `post_get_nas_trial_detail_with_metadata` interceptor runs after the + `post_get_nas_trial_detail` interceptor. The (possibly modified) response returned by + `post_get_nas_trial_detail` will be passed to + `post_get_nas_trial_detail_with_metadata`. + """ + return response, metadata + def pre_list_batch_prediction_jobs( self, request: job_service.ListBatchPredictionJobsRequest, @@ -929,12 +1387,38 @@ def post_list_batch_prediction_jobs( ) -> job_service.ListBatchPredictionJobsResponse: """Post-rpc interceptor for list_batch_prediction_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_batch_prediction_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_batch_prediction_jobs` interceptor runs + before the `post_list_batch_prediction_jobs_with_metadata` interceptor. """ return response + def post_list_batch_prediction_jobs_with_metadata( + self, + response: job_service.ListBatchPredictionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListBatchPredictionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_batch_prediction_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_batch_prediction_jobs_with_metadata` + interceptor in new development instead of the `post_list_batch_prediction_jobs` interceptor. + When both interceptors are used, this `post_list_batch_prediction_jobs_with_metadata` interceptor runs after the + `post_list_batch_prediction_jobs` interceptor. The (possibly modified) response returned by + `post_list_batch_prediction_jobs` will be passed to + `post_list_batch_prediction_jobs_with_metadata`. + """ + return response, metadata + def pre_list_custom_jobs( self, request: job_service.ListCustomJobsRequest, @@ -954,12 +1438,37 @@ def post_list_custom_jobs( ) -> job_service.ListCustomJobsResponse: """Post-rpc interceptor for list_custom_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_custom_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_custom_jobs` interceptor runs + before the `post_list_custom_jobs_with_metadata` interceptor. """ return response + def post_list_custom_jobs_with_metadata( + self, + response: job_service.ListCustomJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListCustomJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_custom_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_custom_jobs_with_metadata` + interceptor in new development instead of the `post_list_custom_jobs` interceptor. + When both interceptors are used, this `post_list_custom_jobs_with_metadata` interceptor runs after the + `post_list_custom_jobs` interceptor. The (possibly modified) response returned by + `post_list_custom_jobs` will be passed to + `post_list_custom_jobs_with_metadata`. + """ + return response, metadata + def pre_list_data_labeling_jobs( self, request: job_service.ListDataLabelingJobsRequest, @@ -979,12 +1488,38 @@ def post_list_data_labeling_jobs( ) -> job_service.ListDataLabelingJobsResponse: """Post-rpc interceptor for list_data_labeling_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_labeling_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_labeling_jobs` interceptor runs + before the `post_list_data_labeling_jobs_with_metadata` interceptor. """ return response + def post_list_data_labeling_jobs_with_metadata( + self, + response: job_service.ListDataLabelingJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListDataLabelingJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_data_labeling_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_data_labeling_jobs_with_metadata` + interceptor in new development instead of the `post_list_data_labeling_jobs` interceptor. + When both interceptors are used, this `post_list_data_labeling_jobs_with_metadata` interceptor runs after the + `post_list_data_labeling_jobs` interceptor. The (possibly modified) response returned by + `post_list_data_labeling_jobs` will be passed to + `post_list_data_labeling_jobs_with_metadata`. + """ + return response, metadata + def pre_list_hyperparameter_tuning_jobs( self, request: job_service.ListHyperparameterTuningJobsRequest, @@ -1005,12 +1540,38 @@ def post_list_hyperparameter_tuning_jobs( ) -> job_service.ListHyperparameterTuningJobsResponse: """Post-rpc interceptor for list_hyperparameter_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_hyperparameter_tuning_jobs` interceptor runs + before the `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor. """ return response + def post_list_hyperparameter_tuning_jobs_with_metadata( + self, + response: job_service.ListHyperparameterTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListHyperparameterTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_hyperparameter_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_hyperparameter_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor runs after the + `post_list_hyperparameter_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_hyperparameter_tuning_jobs` will be passed to + `post_list_hyperparameter_tuning_jobs_with_metadata`. + """ + return response, metadata + def pre_list_model_deployment_monitoring_jobs( self, request: job_service.ListModelDeploymentMonitoringJobsRequest, @@ -1031,12 +1592,38 @@ def post_list_model_deployment_monitoring_jobs( ) -> job_service.ListModelDeploymentMonitoringJobsResponse: """Post-rpc interceptor for list_model_deployment_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_deployment_monitoring_jobs` interceptor runs + before the `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor. """ return response + def post_list_model_deployment_monitoring_jobs_with_metadata( + self, + response: job_service.ListModelDeploymentMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListModelDeploymentMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_deployment_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_deployment_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_deployment_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_deployment_monitoring_jobs` will be passed to + `post_list_model_deployment_monitoring_jobs_with_metadata`. + """ + return response, metadata + def pre_list_nas_jobs( self, request: job_service.ListNasJobsRequest, @@ -1054,12 +1641,37 @@ def post_list_nas_jobs( ) -> job_service.ListNasJobsResponse: """Post-rpc interceptor for list_nas_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_jobs` interceptor runs + before the `post_list_nas_jobs_with_metadata` interceptor. """ return response + def post_list_nas_jobs_with_metadata( + self, + response: job_service.ListNasJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_jobs_with_metadata` + interceptor in new development instead of the `post_list_nas_jobs` interceptor. + When both interceptors are used, this `post_list_nas_jobs_with_metadata` interceptor runs after the + `post_list_nas_jobs` interceptor. The (possibly modified) response returned by + `post_list_nas_jobs` will be passed to + `post_list_nas_jobs_with_metadata`. + """ + return response, metadata + def pre_list_nas_trial_details( self, request: job_service.ListNasTrialDetailsRequest, @@ -1079,12 +1691,37 @@ def post_list_nas_trial_details( ) -> job_service.ListNasTrialDetailsResponse: """Post-rpc interceptor for list_nas_trial_details - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_trial_details_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_trial_details` interceptor runs + before the `post_list_nas_trial_details_with_metadata` interceptor. """ return response + def post_list_nas_trial_details_with_metadata( + self, + response: job_service.ListNasTrialDetailsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasTrialDetailsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_trial_details + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_trial_details_with_metadata` + interceptor in new development instead of the `post_list_nas_trial_details` interceptor. + When both interceptors are used, this `post_list_nas_trial_details_with_metadata` interceptor runs after the + `post_list_nas_trial_details` interceptor. The (possibly modified) response returned by + `post_list_nas_trial_details` will be passed to + `post_list_nas_trial_details_with_metadata`. + """ + return response, metadata + def pre_pause_model_deployment_monitoring_job( self, request: job_service.PauseModelDeploymentMonitoringJobRequest, @@ -1136,12 +1773,38 @@ def post_search_model_deployment_monitoring_stats_anomalies( ) -> job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse: """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_deployment_monitoring_stats_anomalies` interceptor runs + before the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor. """ return response + def post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + self, + response: job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor in new development instead of the `post_search_model_deployment_monitoring_stats_anomalies` interceptor. + When both interceptors are used, this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor runs after the + `post_search_model_deployment_monitoring_stats_anomalies` interceptor. The (possibly modified) response returned by + `post_search_model_deployment_monitoring_stats_anomalies` will be passed to + `post_search_model_deployment_monitoring_stats_anomalies_with_metadata`. + """ + return response, metadata + def pre_update_model_deployment_monitoring_job( self, request: job_service.UpdateModelDeploymentMonitoringJobRequest, @@ -1162,12 +1825,35 @@ def post_update_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_deployment_monitoring_job` interceptor runs + before the `post_update_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_update_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_update_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_update_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_update_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_update_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_update_model_deployment_monitoring_job` will be passed to + `post_update_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -4137,6 +4823,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4297,6 +4987,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4450,6 +5144,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4610,6 +5308,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4778,6 +5483,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4936,6 +5648,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5082,6 +5798,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5227,6 +5947,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5372,6 +6096,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5521,6 +6249,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5671,6 +6406,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5818,6 +6560,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5969,6 +6715,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6125,6 +6875,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6272,6 +7026,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6425,6 +7183,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6587,6 +7352,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6735,6 +7507,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6883,6 +7659,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_nas_trial_detail(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_nas_trial_detail_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7030,6 +7810,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_batch_prediction_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_batch_prediction_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7180,6 +7964,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_custom_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_custom_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7328,6 +8116,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_data_labeling_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_data_labeling_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7480,6 +8272,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_hyperparameter_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_hyperparameter_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7635,6 +8434,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_deployment_monitoring_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_model_deployment_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7787,6 +8593,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_nas_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_nas_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7933,6 +8743,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_nas_trial_details(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_nas_trial_details_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8342,6 +9156,13 @@ def __call__( resp = self._interceptor.post_search_model_deployment_monitoring_stats_anomalies( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8500,6 +9321,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/job_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/job_service/transports/rest_asyncio.py index 3573a3d67b..fd15c945f6 100644 --- a/google/cloud/aiplatform_v1/services/job_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/job_service/transports/rest_asyncio.py @@ -467,12 +467,38 @@ async def post_create_batch_prediction_job( ) -> gca_batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for create_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_batch_prediction_job` interceptor runs + before the `post_create_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_create_batch_prediction_job_with_metadata( + self, + response: gca_batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_batch_prediction_job.BatchPredictionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_create_batch_prediction_job` interceptor. + When both interceptors are used, this `post_create_batch_prediction_job_with_metadata` interceptor runs after the + `post_create_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_create_batch_prediction_job` will be passed to + `post_create_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_create_custom_job( self, request: job_service.CreateCustomJobRequest, @@ -492,12 +518,35 @@ async def post_create_custom_job( ) -> gca_custom_job.CustomJob: """Post-rpc interceptor for create_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_custom_job` interceptor runs + before the `post_create_custom_job_with_metadata` interceptor. """ return response + async def post_create_custom_job_with_metadata( + self, + response: gca_custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_custom_job_with_metadata` + interceptor in new development instead of the `post_create_custom_job` interceptor. + When both interceptors are used, this `post_create_custom_job_with_metadata` interceptor runs after the + `post_create_custom_job` interceptor. The (possibly modified) response returned by + `post_create_custom_job` will be passed to + `post_create_custom_job_with_metadata`. + """ + return response, metadata + async def pre_create_data_labeling_job( self, request: job_service.CreateDataLabelingJobRequest, @@ -518,12 +567,37 @@ async def post_create_data_labeling_job( ) -> gca_data_labeling_job.DataLabelingJob: """Post-rpc interceptor for create_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_data_labeling_job` interceptor runs + before the `post_create_data_labeling_job_with_metadata` interceptor. """ return response + async def post_create_data_labeling_job_with_metadata( + self, + response: gca_data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_create_data_labeling_job` interceptor. + When both interceptors are used, this `post_create_data_labeling_job_with_metadata` interceptor runs after the + `post_create_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_create_data_labeling_job` will be passed to + `post_create_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_create_hyperparameter_tuning_job( self, request: job_service.CreateHyperparameterTuningJobRequest, @@ -544,12 +618,38 @@ async def post_create_hyperparameter_tuning_job( ) -> gca_hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for create_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_hyperparameter_tuning_job` interceptor runs + before the `post_create_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_create_hyperparameter_tuning_job_with_metadata( + self, + response: gca_hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_create_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_create_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_hyperparameter_tuning_job` will be passed to + `post_create_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_create_model_deployment_monitoring_job( self, request: job_service.CreateModelDeploymentMonitoringJobRequest, @@ -570,12 +670,38 @@ async def post_create_model_deployment_monitoring_job( ) -> gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for create_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_deployment_monitoring_job` interceptor runs + before the `post_create_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_create_model_deployment_monitoring_job_with_metadata( + self, + response: gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_deployment_monitoring_job` will be passed to + `post_create_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_create_nas_job( self, request: job_service.CreateNasJobRequest, @@ -595,12 +721,35 @@ async def post_create_nas_job( ) -> gca_nas_job.NasJob: """Post-rpc interceptor for create_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_nas_job` interceptor runs + before the `post_create_nas_job_with_metadata` interceptor. """ return response + async def post_create_nas_job_with_metadata( + self, + response: gca_nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_nas_job_with_metadata` + interceptor in new development instead of the `post_create_nas_job` interceptor. + When both interceptors are used, this `post_create_nas_job_with_metadata` interceptor runs after the + `post_create_nas_job` interceptor. The (possibly modified) response returned by + `post_create_nas_job` will be passed to + `post_create_nas_job_with_metadata`. + """ + return response, metadata + async def pre_delete_batch_prediction_job( self, request: job_service.DeleteBatchPredictionJobRequest, @@ -621,12 +770,35 @@ async def post_delete_batch_prediction_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_batch_prediction_job` interceptor runs + before the `post_delete_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_delete_batch_prediction_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_delete_batch_prediction_job` interceptor. + When both interceptors are used, this `post_delete_batch_prediction_job_with_metadata` interceptor runs after the + `post_delete_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_delete_batch_prediction_job` will be passed to + `post_delete_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_delete_custom_job( self, request: job_service.DeleteCustomJobRequest, @@ -646,12 +818,35 @@ async def post_delete_custom_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_custom_job` interceptor runs + before the `post_delete_custom_job_with_metadata` interceptor. """ return response + async def post_delete_custom_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_custom_job_with_metadata` + interceptor in new development instead of the `post_delete_custom_job` interceptor. + When both interceptors are used, this `post_delete_custom_job_with_metadata` interceptor runs after the + `post_delete_custom_job` interceptor. The (possibly modified) response returned by + `post_delete_custom_job` will be passed to + `post_delete_custom_job_with_metadata`. + """ + return response, metadata + async def pre_delete_data_labeling_job( self, request: job_service.DeleteDataLabelingJobRequest, @@ -672,12 +867,35 @@ async def post_delete_data_labeling_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_data_labeling_job` interceptor runs + before the `post_delete_data_labeling_job_with_metadata` interceptor. """ return response + async def post_delete_data_labeling_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_delete_data_labeling_job` interceptor. + When both interceptors are used, this `post_delete_data_labeling_job_with_metadata` interceptor runs after the + `post_delete_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_delete_data_labeling_job` will be passed to + `post_delete_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_delete_hyperparameter_tuning_job( self, request: job_service.DeleteHyperparameterTuningJobRequest, @@ -698,12 +916,35 @@ async def post_delete_hyperparameter_tuning_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_hyperparameter_tuning_job` interceptor runs + before the `post_delete_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_delete_hyperparameter_tuning_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_delete_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_delete_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_delete_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_delete_hyperparameter_tuning_job` will be passed to + `post_delete_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_delete_model_deployment_monitoring_job( self, request: job_service.DeleteModelDeploymentMonitoringJobRequest, @@ -724,12 +965,35 @@ async def post_delete_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_deployment_monitoring_job` interceptor runs + before the `post_delete_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_delete_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_deployment_monitoring_job` will be passed to + `post_delete_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_delete_nas_job( self, request: job_service.DeleteNasJobRequest, @@ -749,12 +1013,35 @@ async def post_delete_nas_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_nas_job` interceptor runs + before the `post_delete_nas_job_with_metadata` interceptor. """ return response + async def post_delete_nas_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_nas_job_with_metadata` + interceptor in new development instead of the `post_delete_nas_job` interceptor. + When both interceptors are used, this `post_delete_nas_job_with_metadata` interceptor runs after the + `post_delete_nas_job` interceptor. The (possibly modified) response returned by + `post_delete_nas_job` will be passed to + `post_delete_nas_job_with_metadata`. + """ + return response, metadata + async def pre_get_batch_prediction_job( self, request: job_service.GetBatchPredictionJobRequest, @@ -775,12 +1062,37 @@ async def post_get_batch_prediction_job( ) -> batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for get_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_batch_prediction_job` interceptor runs + before the `post_get_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_get_batch_prediction_job_with_metadata( + self, + response: batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + batch_prediction_job.BatchPredictionJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_get_batch_prediction_job` interceptor. + When both interceptors are used, this `post_get_batch_prediction_job_with_metadata` interceptor runs after the + `post_get_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_get_batch_prediction_job` will be passed to + `post_get_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_get_custom_job( self, request: job_service.GetCustomJobRequest, @@ -800,12 +1112,35 @@ async def post_get_custom_job( ) -> custom_job.CustomJob: """Post-rpc interceptor for get_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_custom_job` interceptor runs + before the `post_get_custom_job_with_metadata` interceptor. """ return response + async def post_get_custom_job_with_metadata( + self, + response: custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_custom_job_with_metadata` + interceptor in new development instead of the `post_get_custom_job` interceptor. + When both interceptors are used, this `post_get_custom_job_with_metadata` interceptor runs after the + `post_get_custom_job` interceptor. The (possibly modified) response returned by + `post_get_custom_job` will be passed to + `post_get_custom_job_with_metadata`. + """ + return response, metadata + async def pre_get_data_labeling_job( self, request: job_service.GetDataLabelingJobRequest, @@ -825,12 +1160,37 @@ async def post_get_data_labeling_job( ) -> data_labeling_job.DataLabelingJob: """Post-rpc interceptor for get_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_data_labeling_job` interceptor runs + before the `post_get_data_labeling_job_with_metadata` interceptor. """ return response + async def post_get_data_labeling_job_with_metadata( + self, + response: data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_get_data_labeling_job` interceptor. + When both interceptors are used, this `post_get_data_labeling_job_with_metadata` interceptor runs after the + `post_get_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_get_data_labeling_job` will be passed to + `post_get_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_get_hyperparameter_tuning_job( self, request: job_service.GetHyperparameterTuningJobRequest, @@ -851,12 +1211,38 @@ async def post_get_hyperparameter_tuning_job( ) -> hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for get_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_hyperparameter_tuning_job` interceptor runs + before the `post_get_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_get_hyperparameter_tuning_job_with_metadata( + self, + response: hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_get_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_get_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_hyperparameter_tuning_job` will be passed to + `post_get_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_get_model_deployment_monitoring_job( self, request: job_service.GetModelDeploymentMonitoringJobRequest, @@ -877,12 +1263,38 @@ async def post_get_model_deployment_monitoring_job( ) -> model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for get_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_deployment_monitoring_job` interceptor runs + before the `post_get_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_get_model_deployment_monitoring_job_with_metadata( + self, + response: model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_deployment_monitoring_job` will be passed to + `post_get_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_get_nas_job( self, request: job_service.GetNasJobRequest, @@ -898,12 +1310,35 @@ async def pre_get_nas_job( async def post_get_nas_job(self, response: nas_job.NasJob) -> nas_job.NasJob: """Post-rpc interceptor for get_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_job` interceptor runs + before the `post_get_nas_job_with_metadata` interceptor. """ return response + async def post_get_nas_job_with_metadata( + self, + response: nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_job_with_metadata` + interceptor in new development instead of the `post_get_nas_job` interceptor. + When both interceptors are used, this `post_get_nas_job_with_metadata` interceptor runs after the + `post_get_nas_job` interceptor. The (possibly modified) response returned by + `post_get_nas_job` will be passed to + `post_get_nas_job_with_metadata`. + """ + return response, metadata + async def pre_get_nas_trial_detail( self, request: job_service.GetNasTrialDetailRequest, @@ -923,12 +1358,35 @@ async def post_get_nas_trial_detail( ) -> nas_job.NasTrialDetail: """Post-rpc interceptor for get_nas_trial_detail - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_trial_detail_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_trial_detail` interceptor runs + before the `post_get_nas_trial_detail_with_metadata` interceptor. """ return response + async def post_get_nas_trial_detail_with_metadata( + self, + response: nas_job.NasTrialDetail, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasTrialDetail, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_trial_detail + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_trial_detail_with_metadata` + interceptor in new development instead of the `post_get_nas_trial_detail` interceptor. + When both interceptors are used, this `post_get_nas_trial_detail_with_metadata` interceptor runs after the + `post_get_nas_trial_detail` interceptor. The (possibly modified) response returned by + `post_get_nas_trial_detail` will be passed to + `post_get_nas_trial_detail_with_metadata`. + """ + return response, metadata + async def pre_list_batch_prediction_jobs( self, request: job_service.ListBatchPredictionJobsRequest, @@ -949,12 +1407,38 @@ async def post_list_batch_prediction_jobs( ) -> job_service.ListBatchPredictionJobsResponse: """Post-rpc interceptor for list_batch_prediction_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_batch_prediction_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_batch_prediction_jobs` interceptor runs + before the `post_list_batch_prediction_jobs_with_metadata` interceptor. """ return response + async def post_list_batch_prediction_jobs_with_metadata( + self, + response: job_service.ListBatchPredictionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListBatchPredictionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_batch_prediction_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_batch_prediction_jobs_with_metadata` + interceptor in new development instead of the `post_list_batch_prediction_jobs` interceptor. + When both interceptors are used, this `post_list_batch_prediction_jobs_with_metadata` interceptor runs after the + `post_list_batch_prediction_jobs` interceptor. The (possibly modified) response returned by + `post_list_batch_prediction_jobs` will be passed to + `post_list_batch_prediction_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_custom_jobs( self, request: job_service.ListCustomJobsRequest, @@ -974,12 +1458,37 @@ async def post_list_custom_jobs( ) -> job_service.ListCustomJobsResponse: """Post-rpc interceptor for list_custom_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_custom_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_custom_jobs` interceptor runs + before the `post_list_custom_jobs_with_metadata` interceptor. """ return response + async def post_list_custom_jobs_with_metadata( + self, + response: job_service.ListCustomJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListCustomJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_custom_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_custom_jobs_with_metadata` + interceptor in new development instead of the `post_list_custom_jobs` interceptor. + When both interceptors are used, this `post_list_custom_jobs_with_metadata` interceptor runs after the + `post_list_custom_jobs` interceptor. The (possibly modified) response returned by + `post_list_custom_jobs` will be passed to + `post_list_custom_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_data_labeling_jobs( self, request: job_service.ListDataLabelingJobsRequest, @@ -999,12 +1508,38 @@ async def post_list_data_labeling_jobs( ) -> job_service.ListDataLabelingJobsResponse: """Post-rpc interceptor for list_data_labeling_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_labeling_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_labeling_jobs` interceptor runs + before the `post_list_data_labeling_jobs_with_metadata` interceptor. """ return response + async def post_list_data_labeling_jobs_with_metadata( + self, + response: job_service.ListDataLabelingJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListDataLabelingJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_data_labeling_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_data_labeling_jobs_with_metadata` + interceptor in new development instead of the `post_list_data_labeling_jobs` interceptor. + When both interceptors are used, this `post_list_data_labeling_jobs_with_metadata` interceptor runs after the + `post_list_data_labeling_jobs` interceptor. The (possibly modified) response returned by + `post_list_data_labeling_jobs` will be passed to + `post_list_data_labeling_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_hyperparameter_tuning_jobs( self, request: job_service.ListHyperparameterTuningJobsRequest, @@ -1025,12 +1560,38 @@ async def post_list_hyperparameter_tuning_jobs( ) -> job_service.ListHyperparameterTuningJobsResponse: """Post-rpc interceptor for list_hyperparameter_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_hyperparameter_tuning_jobs` interceptor runs + before the `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor. """ return response + async def post_list_hyperparameter_tuning_jobs_with_metadata( + self, + response: job_service.ListHyperparameterTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListHyperparameterTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_hyperparameter_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_hyperparameter_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor runs after the + `post_list_hyperparameter_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_hyperparameter_tuning_jobs` will be passed to + `post_list_hyperparameter_tuning_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_model_deployment_monitoring_jobs( self, request: job_service.ListModelDeploymentMonitoringJobsRequest, @@ -1051,12 +1612,38 @@ async def post_list_model_deployment_monitoring_jobs( ) -> job_service.ListModelDeploymentMonitoringJobsResponse: """Post-rpc interceptor for list_model_deployment_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_deployment_monitoring_jobs` interceptor runs + before the `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor. """ return response + async def post_list_model_deployment_monitoring_jobs_with_metadata( + self, + response: job_service.ListModelDeploymentMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListModelDeploymentMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_deployment_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_deployment_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_deployment_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_deployment_monitoring_jobs` will be passed to + `post_list_model_deployment_monitoring_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_nas_jobs( self, request: job_service.ListNasJobsRequest, @@ -1074,12 +1661,37 @@ async def post_list_nas_jobs( ) -> job_service.ListNasJobsResponse: """Post-rpc interceptor for list_nas_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_jobs` interceptor runs + before the `post_list_nas_jobs_with_metadata` interceptor. """ return response + async def post_list_nas_jobs_with_metadata( + self, + response: job_service.ListNasJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_jobs_with_metadata` + interceptor in new development instead of the `post_list_nas_jobs` interceptor. + When both interceptors are used, this `post_list_nas_jobs_with_metadata` interceptor runs after the + `post_list_nas_jobs` interceptor. The (possibly modified) response returned by + `post_list_nas_jobs` will be passed to + `post_list_nas_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_nas_trial_details( self, request: job_service.ListNasTrialDetailsRequest, @@ -1099,12 +1711,37 @@ async def post_list_nas_trial_details( ) -> job_service.ListNasTrialDetailsResponse: """Post-rpc interceptor for list_nas_trial_details - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_trial_details_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_trial_details` interceptor runs + before the `post_list_nas_trial_details_with_metadata` interceptor. """ return response + async def post_list_nas_trial_details_with_metadata( + self, + response: job_service.ListNasTrialDetailsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasTrialDetailsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_trial_details + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_trial_details_with_metadata` + interceptor in new development instead of the `post_list_nas_trial_details` interceptor. + When both interceptors are used, this `post_list_nas_trial_details_with_metadata` interceptor runs after the + `post_list_nas_trial_details` interceptor. The (possibly modified) response returned by + `post_list_nas_trial_details` will be passed to + `post_list_nas_trial_details_with_metadata`. + """ + return response, metadata + async def pre_pause_model_deployment_monitoring_job( self, request: job_service.PauseModelDeploymentMonitoringJobRequest, @@ -1156,12 +1793,38 @@ async def post_search_model_deployment_monitoring_stats_anomalies( ) -> job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse: """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_deployment_monitoring_stats_anomalies` interceptor runs + before the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor. """ return response + async def post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + self, + response: job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor in new development instead of the `post_search_model_deployment_monitoring_stats_anomalies` interceptor. + When both interceptors are used, this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor runs after the + `post_search_model_deployment_monitoring_stats_anomalies` interceptor. The (possibly modified) response returned by + `post_search_model_deployment_monitoring_stats_anomalies` will be passed to + `post_search_model_deployment_monitoring_stats_anomalies_with_metadata`. + """ + return response, metadata + async def pre_update_model_deployment_monitoring_job( self, request: job_service.UpdateModelDeploymentMonitoringJobRequest, @@ -1182,12 +1845,35 @@ async def post_update_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_deployment_monitoring_job` interceptor runs + before the `post_update_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_update_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_update_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_update_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_update_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_update_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_update_model_deployment_monitoring_job` will be passed to + `post_update_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2503,6 +3189,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2671,6 +3364,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2832,6 +3529,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2999,6 +3703,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3177,6 +3888,13 @@ async def __call__( resp = await self._interceptor.post_create_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3343,6 +4061,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3498,6 +4220,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3653,6 +4382,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3808,6 +4541,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3966,6 +4706,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4128,6 +4875,13 @@ async def __call__( resp = await self._interceptor.post_delete_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4285,6 +5039,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4444,6 +5202,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4608,6 +5373,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4763,6 +5532,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4924,6 +5697,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5092,6 +5872,13 @@ async def __call__( resp = await self._interceptor.post_get_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5250,6 +6037,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5406,6 +6197,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_nas_trial_detail(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_nas_trial_detail_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5560,6 +6355,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_batch_prediction_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_batch_prediction_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5716,6 +6518,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_custom_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_custom_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5872,6 +6678,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_data_labeling_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_data_labeling_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6031,6 +6844,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_hyperparameter_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_hyperparameter_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6196,6 +7016,13 @@ async def __call__( resp = await self._interceptor.post_list_model_deployment_monitoring_jobs( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_deployment_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6356,6 +7183,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_nas_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_nas_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6510,6 +7341,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_nas_trial_details(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_nas_trial_details_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6941,6 +7776,13 @@ async def __call__( resp = await self._interceptor.post_search_model_deployment_monitoring_stats_anomalies( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7111,6 +7953,13 @@ async def __call__( resp = await self._interceptor.post_update_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/llm_utility_service/client.py b/google/cloud/aiplatform_v1/services/llm_utility_service/client.py index fdf78309a8..59a29dbfa8 100644 --- a/google/cloud/aiplatform_v1/services/llm_utility_service/client.py +++ b/google/cloud/aiplatform_v1/services/llm_utility_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -525,6 +527,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1043,16 +1072,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1098,16 +1131,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1270,16 +1307,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1391,16 +1432,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1513,16 +1558,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1573,16 +1622,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1628,16 +1681,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1683,16 +1740,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest.py b/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest.py index 2f2e854c77..93e4ee11ff 100644 --- a/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest.py @@ -121,12 +121,38 @@ def post_compute_tokens( ) -> llm_utility_service.ComputeTokensResponse: """Post-rpc interceptor for compute_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_compute_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_compute_tokens` interceptor runs + before the `post_compute_tokens_with_metadata` interceptor. """ return response + def post_compute_tokens_with_metadata( + self, + response: llm_utility_service.ComputeTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + llm_utility_service.ComputeTokensResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for compute_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_compute_tokens_with_metadata` + interceptor in new development instead of the `post_compute_tokens` interceptor. + When both interceptors are used, this `post_compute_tokens_with_metadata` interceptor runs after the + `post_compute_tokens` interceptor. The (possibly modified) response returned by + `post_compute_tokens` will be passed to + `post_compute_tokens_with_metadata`. + """ + return response, metadata + def pre_count_tokens( self, request: prediction_service.CountTokensRequest, @@ -146,12 +172,37 @@ def post_count_tokens( ) -> prediction_service.CountTokensResponse: """Post-rpc interceptor for count_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_count_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_count_tokens` interceptor runs + before the `post_count_tokens_with_metadata` interceptor. """ return response + def post_count_tokens_with_metadata( + self, + response: prediction_service.CountTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.CountTokensResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for count_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_count_tokens_with_metadata` + interceptor in new development instead of the `post_count_tokens` interceptor. + When both interceptors are used, this `post_count_tokens_with_metadata` interceptor runs after the + `post_count_tokens` interceptor. The (possibly modified) response returned by + `post_count_tokens` will be passed to + `post_count_tokens_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -610,6 +661,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_compute_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_compute_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -759,6 +814,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_count_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_count_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest_asyncio.py index 4d1c7daaeb..425af9c292 100644 --- a/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/llm_utility_service/transports/rest_asyncio.py @@ -138,12 +138,38 @@ async def post_compute_tokens( ) -> llm_utility_service.ComputeTokensResponse: """Post-rpc interceptor for compute_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_compute_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_compute_tokens` interceptor runs + before the `post_compute_tokens_with_metadata` interceptor. """ return response + async def post_compute_tokens_with_metadata( + self, + response: llm_utility_service.ComputeTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + llm_utility_service.ComputeTokensResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for compute_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_compute_tokens_with_metadata` + interceptor in new development instead of the `post_compute_tokens` interceptor. + When both interceptors are used, this `post_compute_tokens_with_metadata` interceptor runs after the + `post_compute_tokens` interceptor. The (possibly modified) response returned by + `post_compute_tokens` will be passed to + `post_compute_tokens_with_metadata`. + """ + return response, metadata + async def pre_count_tokens( self, request: prediction_service.CountTokensRequest, @@ -163,12 +189,37 @@ async def post_count_tokens( ) -> prediction_service.CountTokensResponse: """Post-rpc interceptor for count_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_count_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_count_tokens` interceptor runs + before the `post_count_tokens_with_metadata` interceptor. """ return response + async def post_count_tokens_with_metadata( + self, + response: prediction_service.CountTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.CountTokensResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for count_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_count_tokens_with_metadata` + interceptor in new development instead of the `post_count_tokens` interceptor. + When both interceptors are used, this `post_count_tokens_with_metadata` interceptor runs after the + `post_count_tokens` interceptor. The (possibly modified) response returned by + `post_count_tokens` will be passed to + `post_count_tokens_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -685,6 +736,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_compute_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_compute_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -845,6 +900,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_count_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_count_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/match_service/client.py b/google/cloud/aiplatform_v1/services/match_service/client.py index 49d424c2b6..90fb43699c 100644 --- a/google/cloud/aiplatform_v1/services/match_service/client.py +++ b/google/cloud/aiplatform_v1/services/match_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -501,6 +503,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -951,16 +980,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1006,16 +1039,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1178,16 +1215,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1299,16 +1340,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1421,16 +1466,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1481,16 +1530,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1536,16 +1589,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1591,16 +1648,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/match_service/transports/rest.py b/google/cloud/aiplatform_v1/services/match_service/transports/rest.py index 7ee3ea6e99..08988a7c40 100644 --- a/google/cloud/aiplatform_v1/services/match_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/match_service/transports/rest.py @@ -119,12 +119,37 @@ def post_find_neighbors( ) -> match_service.FindNeighborsResponse: """Post-rpc interceptor for find_neighbors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_find_neighbors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_find_neighbors` interceptor runs + before the `post_find_neighbors_with_metadata` interceptor. """ return response + def post_find_neighbors_with_metadata( + self, + response: match_service.FindNeighborsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.FindNeighborsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for find_neighbors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_find_neighbors_with_metadata` + interceptor in new development instead of the `post_find_neighbors` interceptor. + When both interceptors are used, this `post_find_neighbors_with_metadata` interceptor runs after the + `post_find_neighbors` interceptor. The (possibly modified) response returned by + `post_find_neighbors` will be passed to + `post_find_neighbors_with_metadata`. + """ + return response, metadata + def pre_read_index_datapoints( self, request: match_service.ReadIndexDatapointsRequest, @@ -145,12 +170,38 @@ def post_read_index_datapoints( ) -> match_service.ReadIndexDatapointsResponse: """Post-rpc interceptor for read_index_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_index_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_read_index_datapoints` interceptor runs + before the `post_read_index_datapoints_with_metadata` interceptor. """ return response + def post_read_index_datapoints_with_metadata( + self, + response: match_service.ReadIndexDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.ReadIndexDatapointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_index_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_read_index_datapoints_with_metadata` + interceptor in new development instead of the `post_read_index_datapoints` interceptor. + When both interceptors are used, this `post_read_index_datapoints_with_metadata` interceptor runs after the + `post_read_index_datapoints` interceptor. The (possibly modified) response returned by + `post_read_index_datapoints` will be passed to + `post_read_index_datapoints_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -609,6 +660,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_find_neighbors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_find_neighbors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -763,6 +818,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_index_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_index_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/match_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/match_service/transports/rest_asyncio.py index 5a1dbd3c4c..85d4d2f9dd 100644 --- a/google/cloud/aiplatform_v1/services/match_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/match_service/transports/rest_asyncio.py @@ -136,12 +136,37 @@ async def post_find_neighbors( ) -> match_service.FindNeighborsResponse: """Post-rpc interceptor for find_neighbors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_find_neighbors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_find_neighbors` interceptor runs + before the `post_find_neighbors_with_metadata` interceptor. """ return response + async def post_find_neighbors_with_metadata( + self, + response: match_service.FindNeighborsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.FindNeighborsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for find_neighbors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_find_neighbors_with_metadata` + interceptor in new development instead of the `post_find_neighbors` interceptor. + When both interceptors are used, this `post_find_neighbors_with_metadata` interceptor runs after the + `post_find_neighbors` interceptor. The (possibly modified) response returned by + `post_find_neighbors` will be passed to + `post_find_neighbors_with_metadata`. + """ + return response, metadata + async def pre_read_index_datapoints( self, request: match_service.ReadIndexDatapointsRequest, @@ -162,12 +187,38 @@ async def post_read_index_datapoints( ) -> match_service.ReadIndexDatapointsResponse: """Post-rpc interceptor for read_index_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_index_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_read_index_datapoints` interceptor runs + before the `post_read_index_datapoints_with_metadata` interceptor. """ return response + async def post_read_index_datapoints_with_metadata( + self, + response: match_service.ReadIndexDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.ReadIndexDatapointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_index_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_read_index_datapoints_with_metadata` + interceptor in new development instead of the `post_read_index_datapoints` interceptor. + When both interceptors are used, this `post_read_index_datapoints_with_metadata` interceptor runs after the + `post_read_index_datapoints` interceptor. The (possibly modified) response returned by + `post_read_index_datapoints` will be passed to + `post_read_index_datapoints_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -684,6 +735,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_find_neighbors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_find_neighbors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -847,6 +902,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_index_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_index_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/metadata_service/client.py b/google/cloud/aiplatform_v1/services/metadata_service/client.py index 44a0ec4e44..f3fd6e45c6 100644 --- a/google/cloud/aiplatform_v1/services/metadata_service/client.py +++ b/google/cloud/aiplatform_v1/services/metadata_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -617,6 +619,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -4805,16 +4834,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -4860,16 +4893,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -5032,16 +5069,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5153,16 +5194,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5275,16 +5320,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5335,16 +5384,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -5390,16 +5443,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -5445,16 +5502,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/metadata_service/transports/rest.py b/google/cloud/aiplatform_v1/services/metadata_service/transports/rest.py index 1718471036..c8ad218f3b 100644 --- a/google/cloud/aiplatform_v1/services/metadata_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/metadata_service/transports/rest.py @@ -371,12 +371,38 @@ def post_add_context_artifacts_and_executions( ) -> metadata_service.AddContextArtifactsAndExecutionsResponse: """Post-rpc interceptor for add_context_artifacts_and_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_artifacts_and_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_artifacts_and_executions` interceptor runs + before the `post_add_context_artifacts_and_executions_with_metadata` interceptor. """ return response + def post_add_context_artifacts_and_executions_with_metadata( + self, + response: metadata_service.AddContextArtifactsAndExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextArtifactsAndExecutionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_artifacts_and_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_artifacts_and_executions_with_metadata` + interceptor in new development instead of the `post_add_context_artifacts_and_executions` interceptor. + When both interceptors are used, this `post_add_context_artifacts_and_executions_with_metadata` interceptor runs after the + `post_add_context_artifacts_and_executions` interceptor. The (possibly modified) response returned by + `post_add_context_artifacts_and_executions` will be passed to + `post_add_context_artifacts_and_executions_with_metadata`. + """ + return response, metadata + def pre_add_context_children( self, request: metadata_service.AddContextChildrenRequest, @@ -397,12 +423,38 @@ def post_add_context_children( ) -> metadata_service.AddContextChildrenResponse: """Post-rpc interceptor for add_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_children` interceptor runs + before the `post_add_context_children_with_metadata` interceptor. """ return response + def post_add_context_children_with_metadata( + self, + response: metadata_service.AddContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_children_with_metadata` + interceptor in new development instead of the `post_add_context_children` interceptor. + When both interceptors are used, this `post_add_context_children_with_metadata` interceptor runs after the + `post_add_context_children` interceptor. The (possibly modified) response returned by + `post_add_context_children` will be passed to + `post_add_context_children_with_metadata`. + """ + return response, metadata + def pre_add_execution_events( self, request: metadata_service.AddExecutionEventsRequest, @@ -423,12 +475,38 @@ def post_add_execution_events( ) -> metadata_service.AddExecutionEventsResponse: """Post-rpc interceptor for add_execution_events - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_execution_events_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_execution_events` interceptor runs + before the `post_add_execution_events_with_metadata` interceptor. """ return response + def post_add_execution_events_with_metadata( + self, + response: metadata_service.AddExecutionEventsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddExecutionEventsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_execution_events + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_execution_events_with_metadata` + interceptor in new development instead of the `post_add_execution_events` interceptor. + When both interceptors are used, this `post_add_execution_events_with_metadata` interceptor runs after the + `post_add_execution_events` interceptor. The (possibly modified) response returned by + `post_add_execution_events` will be passed to + `post_add_execution_events_with_metadata`. + """ + return response, metadata + def pre_create_artifact( self, request: metadata_service.CreateArtifactRequest, @@ -448,12 +526,35 @@ def post_create_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for create_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_artifact` interceptor runs + before the `post_create_artifact_with_metadata` interceptor. """ return response + def post_create_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_artifact_with_metadata` + interceptor in new development instead of the `post_create_artifact` interceptor. + When both interceptors are used, this `post_create_artifact_with_metadata` interceptor runs after the + `post_create_artifact` interceptor. The (possibly modified) response returned by + `post_create_artifact` will be passed to + `post_create_artifact_with_metadata`. + """ + return response, metadata + def pre_create_context( self, request: metadata_service.CreateContextRequest, @@ -471,12 +572,35 @@ def pre_create_context( def post_create_context(self, response: gca_context.Context) -> gca_context.Context: """Post-rpc interceptor for create_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_context` interceptor runs + before the `post_create_context_with_metadata` interceptor. """ return response + def post_create_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_context_with_metadata` + interceptor in new development instead of the `post_create_context` interceptor. + When both interceptors are used, this `post_create_context_with_metadata` interceptor runs after the + `post_create_context` interceptor. The (possibly modified) response returned by + `post_create_context` will be passed to + `post_create_context_with_metadata`. + """ + return response, metadata + def pre_create_execution( self, request: metadata_service.CreateExecutionRequest, @@ -496,12 +620,35 @@ def post_create_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for create_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_execution` interceptor runs + before the `post_create_execution_with_metadata` interceptor. """ return response + def post_create_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_execution_with_metadata` + interceptor in new development instead of the `post_create_execution` interceptor. + When both interceptors are used, this `post_create_execution_with_metadata` interceptor runs after the + `post_create_execution` interceptor. The (possibly modified) response returned by + `post_create_execution` will be passed to + `post_create_execution_with_metadata`. + """ + return response, metadata + def pre_create_metadata_schema( self, request: metadata_service.CreateMetadataSchemaRequest, @@ -522,12 +669,37 @@ def post_create_metadata_schema( ) -> gca_metadata_schema.MetadataSchema: """Post-rpc interceptor for create_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_schema` interceptor runs + before the `post_create_metadata_schema_with_metadata` interceptor. """ return response + def post_create_metadata_schema_with_metadata( + self, + response: gca_metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_schema_with_metadata` + interceptor in new development instead of the `post_create_metadata_schema` interceptor. + When both interceptors are used, this `post_create_metadata_schema_with_metadata` interceptor runs after the + `post_create_metadata_schema` interceptor. The (possibly modified) response returned by + `post_create_metadata_schema` will be passed to + `post_create_metadata_schema_with_metadata`. + """ + return response, metadata + def pre_create_metadata_store( self, request: metadata_service.CreateMetadataStoreRequest, @@ -548,12 +720,35 @@ def post_create_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_store` interceptor runs + before the `post_create_metadata_store_with_metadata` interceptor. """ return response + def post_create_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_store_with_metadata` + interceptor in new development instead of the `post_create_metadata_store` interceptor. + When both interceptors are used, this `post_create_metadata_store_with_metadata` interceptor runs after the + `post_create_metadata_store` interceptor. The (possibly modified) response returned by + `post_create_metadata_store` will be passed to + `post_create_metadata_store_with_metadata`. + """ + return response, metadata + def pre_delete_artifact( self, request: metadata_service.DeleteArtifactRequest, @@ -573,12 +768,35 @@ def post_delete_artifact( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_artifact` interceptor runs + before the `post_delete_artifact_with_metadata` interceptor. """ return response + def post_delete_artifact_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_artifact_with_metadata` + interceptor in new development instead of the `post_delete_artifact` interceptor. + When both interceptors are used, this `post_delete_artifact_with_metadata` interceptor runs after the + `post_delete_artifact` interceptor. The (possibly modified) response returned by + `post_delete_artifact` will be passed to + `post_delete_artifact_with_metadata`. + """ + return response, metadata + def pre_delete_context( self, request: metadata_service.DeleteContextRequest, @@ -598,12 +816,35 @@ def post_delete_context( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_context` interceptor runs + before the `post_delete_context_with_metadata` interceptor. """ return response + def post_delete_context_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_context_with_metadata` + interceptor in new development instead of the `post_delete_context` interceptor. + When both interceptors are used, this `post_delete_context_with_metadata` interceptor runs after the + `post_delete_context` interceptor. The (possibly modified) response returned by + `post_delete_context` will be passed to + `post_delete_context_with_metadata`. + """ + return response, metadata + def pre_delete_execution( self, request: metadata_service.DeleteExecutionRequest, @@ -623,12 +864,35 @@ def post_delete_execution( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_execution` interceptor runs + before the `post_delete_execution_with_metadata` interceptor. """ return response + def post_delete_execution_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_execution_with_metadata` + interceptor in new development instead of the `post_delete_execution` interceptor. + When both interceptors are used, this `post_delete_execution_with_metadata` interceptor runs after the + `post_delete_execution` interceptor. The (possibly modified) response returned by + `post_delete_execution` will be passed to + `post_delete_execution_with_metadata`. + """ + return response, metadata + def pre_delete_metadata_store( self, request: metadata_service.DeleteMetadataStoreRequest, @@ -649,12 +913,35 @@ def post_delete_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_metadata_store` interceptor runs + before the `post_delete_metadata_store_with_metadata` interceptor. """ return response + def post_delete_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_metadata_store_with_metadata` + interceptor in new development instead of the `post_delete_metadata_store` interceptor. + When both interceptors are used, this `post_delete_metadata_store_with_metadata` interceptor runs after the + `post_delete_metadata_store` interceptor. The (possibly modified) response returned by + `post_delete_metadata_store` will be passed to + `post_delete_metadata_store_with_metadata`. + """ + return response, metadata + def pre_get_artifact( self, request: metadata_service.GetArtifactRequest, @@ -672,12 +959,35 @@ def pre_get_artifact( def post_get_artifact(self, response: artifact.Artifact) -> artifact.Artifact: """Post-rpc interceptor for get_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_artifact` interceptor runs + before the `post_get_artifact_with_metadata` interceptor. """ return response + def post_get_artifact_with_metadata( + self, + response: artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_artifact_with_metadata` + interceptor in new development instead of the `post_get_artifact` interceptor. + When both interceptors are used, this `post_get_artifact_with_metadata` interceptor runs after the + `post_get_artifact` interceptor. The (possibly modified) response returned by + `post_get_artifact` will be passed to + `post_get_artifact_with_metadata`. + """ + return response, metadata + def pre_get_context( self, request: metadata_service.GetContextRequest, @@ -695,12 +1005,35 @@ def pre_get_context( def post_get_context(self, response: context.Context) -> context.Context: """Post-rpc interceptor for get_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_context` interceptor runs + before the `post_get_context_with_metadata` interceptor. """ return response + def post_get_context_with_metadata( + self, + response: context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_context_with_metadata` + interceptor in new development instead of the `post_get_context` interceptor. + When both interceptors are used, this `post_get_context_with_metadata` interceptor runs after the + `post_get_context` interceptor. The (possibly modified) response returned by + `post_get_context` will be passed to + `post_get_context_with_metadata`. + """ + return response, metadata + def pre_get_execution( self, request: metadata_service.GetExecutionRequest, @@ -718,12 +1051,35 @@ def pre_get_execution( def post_get_execution(self, response: execution.Execution) -> execution.Execution: """Post-rpc interceptor for get_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_execution` interceptor runs + before the `post_get_execution_with_metadata` interceptor. """ return response + def post_get_execution_with_metadata( + self, + response: execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_execution_with_metadata` + interceptor in new development instead of the `post_get_execution` interceptor. + When both interceptors are used, this `post_get_execution_with_metadata` interceptor runs after the + `post_get_execution` interceptor. The (possibly modified) response returned by + `post_get_execution` will be passed to + `post_get_execution_with_metadata`. + """ + return response, metadata + def pre_get_metadata_schema( self, request: metadata_service.GetMetadataSchemaRequest, @@ -744,12 +1100,35 @@ def post_get_metadata_schema( ) -> metadata_schema.MetadataSchema: """Post-rpc interceptor for get_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_schema` interceptor runs + before the `post_get_metadata_schema_with_metadata` interceptor. """ return response + def post_get_metadata_schema_with_metadata( + self, + response: metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_schema_with_metadata` + interceptor in new development instead of the `post_get_metadata_schema` interceptor. + When both interceptors are used, this `post_get_metadata_schema_with_metadata` interceptor runs after the + `post_get_metadata_schema` interceptor. The (possibly modified) response returned by + `post_get_metadata_schema` will be passed to + `post_get_metadata_schema_with_metadata`. + """ + return response, metadata + def pre_get_metadata_store( self, request: metadata_service.GetMetadataStoreRequest, @@ -770,12 +1149,35 @@ def post_get_metadata_store( ) -> metadata_store.MetadataStore: """Post-rpc interceptor for get_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_store` interceptor runs + before the `post_get_metadata_store_with_metadata` interceptor. """ return response + def post_get_metadata_store_with_metadata( + self, + response: metadata_store.MetadataStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_store.MetadataStore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_store_with_metadata` + interceptor in new development instead of the `post_get_metadata_store` interceptor. + When both interceptors are used, this `post_get_metadata_store_with_metadata` interceptor runs after the + `post_get_metadata_store` interceptor. The (possibly modified) response returned by + `post_get_metadata_store` will be passed to + `post_get_metadata_store_with_metadata`. + """ + return response, metadata + def pre_list_artifacts( self, request: metadata_service.ListArtifactsRequest, @@ -795,12 +1197,37 @@ def post_list_artifacts( ) -> metadata_service.ListArtifactsResponse: """Post-rpc interceptor for list_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_artifacts` interceptor runs + before the `post_list_artifacts_with_metadata` interceptor. """ return response + def post_list_artifacts_with_metadata( + self, + response: metadata_service.ListArtifactsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListArtifactsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_artifacts_with_metadata` + interceptor in new development instead of the `post_list_artifacts` interceptor. + When both interceptors are used, this `post_list_artifacts_with_metadata` interceptor runs after the + `post_list_artifacts` interceptor. The (possibly modified) response returned by + `post_list_artifacts` will be passed to + `post_list_artifacts_with_metadata`. + """ + return response, metadata + def pre_list_contexts( self, request: metadata_service.ListContextsRequest, @@ -820,12 +1247,37 @@ def post_list_contexts( ) -> metadata_service.ListContextsResponse: """Post-rpc interceptor for list_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_contexts` interceptor runs + before the `post_list_contexts_with_metadata` interceptor. """ return response + def post_list_contexts_with_metadata( + self, + response: metadata_service.ListContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListContextsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_contexts_with_metadata` + interceptor in new development instead of the `post_list_contexts` interceptor. + When both interceptors are used, this `post_list_contexts_with_metadata` interceptor runs after the + `post_list_contexts` interceptor. The (possibly modified) response returned by + `post_list_contexts` will be passed to + `post_list_contexts_with_metadata`. + """ + return response, metadata + def pre_list_executions( self, request: metadata_service.ListExecutionsRequest, @@ -845,12 +1297,37 @@ def post_list_executions( ) -> metadata_service.ListExecutionsResponse: """Post-rpc interceptor for list_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_executions` interceptor runs + before the `post_list_executions_with_metadata` interceptor. """ return response + def post_list_executions_with_metadata( + self, + response: metadata_service.ListExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListExecutionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_executions_with_metadata` + interceptor in new development instead of the `post_list_executions` interceptor. + When both interceptors are used, this `post_list_executions_with_metadata` interceptor runs after the + `post_list_executions` interceptor. The (possibly modified) response returned by + `post_list_executions` will be passed to + `post_list_executions_with_metadata`. + """ + return response, metadata + def pre_list_metadata_schemas( self, request: metadata_service.ListMetadataSchemasRequest, @@ -871,12 +1348,38 @@ def post_list_metadata_schemas( ) -> metadata_service.ListMetadataSchemasResponse: """Post-rpc interceptor for list_metadata_schemas - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_schemas_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_schemas` interceptor runs + before the `post_list_metadata_schemas_with_metadata` interceptor. """ return response + def post_list_metadata_schemas_with_metadata( + self, + response: metadata_service.ListMetadataSchemasResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataSchemasResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_schemas + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_schemas_with_metadata` + interceptor in new development instead of the `post_list_metadata_schemas` interceptor. + When both interceptors are used, this `post_list_metadata_schemas_with_metadata` interceptor runs after the + `post_list_metadata_schemas` interceptor. The (possibly modified) response returned by + `post_list_metadata_schemas` will be passed to + `post_list_metadata_schemas_with_metadata`. + """ + return response, metadata + def pre_list_metadata_stores( self, request: metadata_service.ListMetadataStoresRequest, @@ -897,12 +1400,38 @@ def post_list_metadata_stores( ) -> metadata_service.ListMetadataStoresResponse: """Post-rpc interceptor for list_metadata_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_stores` interceptor runs + before the `post_list_metadata_stores_with_metadata` interceptor. """ return response + def post_list_metadata_stores_with_metadata( + self, + response: metadata_service.ListMetadataStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_stores_with_metadata` + interceptor in new development instead of the `post_list_metadata_stores` interceptor. + When both interceptors are used, this `post_list_metadata_stores_with_metadata` interceptor runs after the + `post_list_metadata_stores` interceptor. The (possibly modified) response returned by + `post_list_metadata_stores` will be passed to + `post_list_metadata_stores_with_metadata`. + """ + return response, metadata + def pre_purge_artifacts( self, request: metadata_service.PurgeArtifactsRequest, @@ -922,12 +1451,35 @@ def post_purge_artifacts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_artifacts` interceptor runs + before the `post_purge_artifacts_with_metadata` interceptor. """ return response + def post_purge_artifacts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_artifacts_with_metadata` + interceptor in new development instead of the `post_purge_artifacts` interceptor. + When both interceptors are used, this `post_purge_artifacts_with_metadata` interceptor runs after the + `post_purge_artifacts` interceptor. The (possibly modified) response returned by + `post_purge_artifacts` will be passed to + `post_purge_artifacts_with_metadata`. + """ + return response, metadata + def pre_purge_contexts( self, request: metadata_service.PurgeContextsRequest, @@ -947,12 +1499,35 @@ def post_purge_contexts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_contexts` interceptor runs + before the `post_purge_contexts_with_metadata` interceptor. """ return response + def post_purge_contexts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_contexts_with_metadata` + interceptor in new development instead of the `post_purge_contexts` interceptor. + When both interceptors are used, this `post_purge_contexts_with_metadata` interceptor runs after the + `post_purge_contexts` interceptor. The (possibly modified) response returned by + `post_purge_contexts` will be passed to + `post_purge_contexts_with_metadata`. + """ + return response, metadata + def pre_purge_executions( self, request: metadata_service.PurgeExecutionsRequest, @@ -972,12 +1547,35 @@ def post_purge_executions( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_executions` interceptor runs + before the `post_purge_executions_with_metadata` interceptor. """ return response + def post_purge_executions_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_executions_with_metadata` + interceptor in new development instead of the `post_purge_executions` interceptor. + When both interceptors are used, this `post_purge_executions_with_metadata` interceptor runs after the + `post_purge_executions` interceptor. The (possibly modified) response returned by + `post_purge_executions` will be passed to + `post_purge_executions_with_metadata`. + """ + return response, metadata + def pre_query_artifact_lineage_subgraph( self, request: metadata_service.QueryArtifactLineageSubgraphRequest, @@ -998,12 +1596,37 @@ def post_query_artifact_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_artifact_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_artifact_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_artifact_lineage_subgraph` interceptor runs + before the `post_query_artifact_lineage_subgraph_with_metadata` interceptor. """ return response + def post_query_artifact_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_artifact_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_artifact_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_artifact_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_artifact_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_artifact_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_artifact_lineage_subgraph` will be passed to + `post_query_artifact_lineage_subgraph_with_metadata`. + """ + return response, metadata + def pre_query_context_lineage_subgraph( self, request: metadata_service.QueryContextLineageSubgraphRequest, @@ -1024,12 +1647,37 @@ def post_query_context_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_context_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_context_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_context_lineage_subgraph` interceptor runs + before the `post_query_context_lineage_subgraph_with_metadata` interceptor. """ return response + def post_query_context_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_context_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_context_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_context_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_context_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_context_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_context_lineage_subgraph` will be passed to + `post_query_context_lineage_subgraph_with_metadata`. + """ + return response, metadata + def pre_query_execution_inputs_and_outputs( self, request: metadata_service.QueryExecutionInputsAndOutputsRequest, @@ -1050,12 +1698,37 @@ def post_query_execution_inputs_and_outputs( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_execution_inputs_and_outputs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_execution_inputs_and_outputs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_execution_inputs_and_outputs` interceptor runs + before the `post_query_execution_inputs_and_outputs_with_metadata` interceptor. """ return response + def post_query_execution_inputs_and_outputs_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_execution_inputs_and_outputs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_execution_inputs_and_outputs_with_metadata` + interceptor in new development instead of the `post_query_execution_inputs_and_outputs` interceptor. + When both interceptors are used, this `post_query_execution_inputs_and_outputs_with_metadata` interceptor runs after the + `post_query_execution_inputs_and_outputs` interceptor. The (possibly modified) response returned by + `post_query_execution_inputs_and_outputs` will be passed to + `post_query_execution_inputs_and_outputs_with_metadata`. + """ + return response, metadata + def pre_remove_context_children( self, request: metadata_service.RemoveContextChildrenRequest, @@ -1076,12 +1749,38 @@ def post_remove_context_children( ) -> metadata_service.RemoveContextChildrenResponse: """Post-rpc interceptor for remove_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_remove_context_children` interceptor runs + before the `post_remove_context_children_with_metadata` interceptor. """ return response + def post_remove_context_children_with_metadata( + self, + response: metadata_service.RemoveContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.RemoveContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for remove_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_remove_context_children_with_metadata` + interceptor in new development instead of the `post_remove_context_children` interceptor. + When both interceptors are used, this `post_remove_context_children_with_metadata` interceptor runs after the + `post_remove_context_children` interceptor. The (possibly modified) response returned by + `post_remove_context_children` will be passed to + `post_remove_context_children_with_metadata`. + """ + return response, metadata + def pre_update_artifact( self, request: metadata_service.UpdateArtifactRequest, @@ -1101,12 +1800,35 @@ def post_update_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for update_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_artifact` interceptor runs + before the `post_update_artifact_with_metadata` interceptor. """ return response + def post_update_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_artifact_with_metadata` + interceptor in new development instead of the `post_update_artifact` interceptor. + When both interceptors are used, this `post_update_artifact_with_metadata` interceptor runs after the + `post_update_artifact` interceptor. The (possibly modified) response returned by + `post_update_artifact` will be passed to + `post_update_artifact_with_metadata`. + """ + return response, metadata + def pre_update_context( self, request: metadata_service.UpdateContextRequest, @@ -1124,12 +1846,35 @@ def pre_update_context( def post_update_context(self, response: gca_context.Context) -> gca_context.Context: """Post-rpc interceptor for update_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_context` interceptor runs + before the `post_update_context_with_metadata` interceptor. """ return response + def post_update_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_context_with_metadata` + interceptor in new development instead of the `post_update_context` interceptor. + When both interceptors are used, this `post_update_context_with_metadata` interceptor runs after the + `post_update_context` interceptor. The (possibly modified) response returned by + `post_update_context` will be passed to + `post_update_context_with_metadata`. + """ + return response, metadata + def pre_update_execution( self, request: metadata_service.UpdateExecutionRequest, @@ -1149,12 +1894,35 @@ def post_update_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for update_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_execution` interceptor runs + before the `post_update_execution_with_metadata` interceptor. """ return response + def post_update_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_execution_with_metadata` + interceptor in new development instead of the `post_update_execution` interceptor. + When both interceptors are used, this `post_update_execution_with_metadata` interceptor runs after the + `post_update_execution` interceptor. The (possibly modified) response returned by + `post_update_execution` will be passed to + `post_update_execution_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3534,6 +4302,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_context_artifacts_and_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_add_context_artifacts_and_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3689,6 +4464,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3844,6 +4623,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_execution_events(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_execution_events_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3994,6 +4777,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4142,6 +4929,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4292,6 +5083,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4443,6 +5238,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4597,6 +5396,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4740,6 +5543,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4883,6 +5690,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5028,6 +5839,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5174,6 +5989,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5316,6 +6135,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5458,6 +6281,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5600,6 +6427,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5745,6 +6576,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5892,6 +6727,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6036,6 +6875,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6182,6 +7025,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6328,6 +7175,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6477,6 +7328,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_metadata_schemas(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_metadata_schemas_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6626,6 +7481,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_metadata_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_metadata_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6777,6 +7636,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6926,6 +7789,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7077,6 +7944,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7226,6 +8097,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_artifact_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_artifact_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7379,6 +8257,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_context_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_context_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7533,6 +8418,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_execution_inputs_and_outputs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_execution_inputs_and_outputs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7690,6 +8582,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_remove_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_remove_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7840,6 +8736,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7988,6 +8888,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8138,6 +9042,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/metadata_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/metadata_service/transports/rest_asyncio.py index ec41208648..582398fe1d 100644 --- a/google/cloud/aiplatform_v1/services/metadata_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/metadata_service/transports/rest_asyncio.py @@ -389,12 +389,38 @@ async def post_add_context_artifacts_and_executions( ) -> metadata_service.AddContextArtifactsAndExecutionsResponse: """Post-rpc interceptor for add_context_artifacts_and_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_artifacts_and_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_artifacts_and_executions` interceptor runs + before the `post_add_context_artifacts_and_executions_with_metadata` interceptor. """ return response + async def post_add_context_artifacts_and_executions_with_metadata( + self, + response: metadata_service.AddContextArtifactsAndExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextArtifactsAndExecutionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_artifacts_and_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_artifacts_and_executions_with_metadata` + interceptor in new development instead of the `post_add_context_artifacts_and_executions` interceptor. + When both interceptors are used, this `post_add_context_artifacts_and_executions_with_metadata` interceptor runs after the + `post_add_context_artifacts_and_executions` interceptor. The (possibly modified) response returned by + `post_add_context_artifacts_and_executions` will be passed to + `post_add_context_artifacts_and_executions_with_metadata`. + """ + return response, metadata + async def pre_add_context_children( self, request: metadata_service.AddContextChildrenRequest, @@ -415,12 +441,38 @@ async def post_add_context_children( ) -> metadata_service.AddContextChildrenResponse: """Post-rpc interceptor for add_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_children` interceptor runs + before the `post_add_context_children_with_metadata` interceptor. """ return response + async def post_add_context_children_with_metadata( + self, + response: metadata_service.AddContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_children_with_metadata` + interceptor in new development instead of the `post_add_context_children` interceptor. + When both interceptors are used, this `post_add_context_children_with_metadata` interceptor runs after the + `post_add_context_children` interceptor. The (possibly modified) response returned by + `post_add_context_children` will be passed to + `post_add_context_children_with_metadata`. + """ + return response, metadata + async def pre_add_execution_events( self, request: metadata_service.AddExecutionEventsRequest, @@ -441,12 +493,38 @@ async def post_add_execution_events( ) -> metadata_service.AddExecutionEventsResponse: """Post-rpc interceptor for add_execution_events - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_execution_events_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_execution_events` interceptor runs + before the `post_add_execution_events_with_metadata` interceptor. """ return response + async def post_add_execution_events_with_metadata( + self, + response: metadata_service.AddExecutionEventsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddExecutionEventsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_execution_events + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_execution_events_with_metadata` + interceptor in new development instead of the `post_add_execution_events` interceptor. + When both interceptors are used, this `post_add_execution_events_with_metadata` interceptor runs after the + `post_add_execution_events` interceptor. The (possibly modified) response returned by + `post_add_execution_events` will be passed to + `post_add_execution_events_with_metadata`. + """ + return response, metadata + async def pre_create_artifact( self, request: metadata_service.CreateArtifactRequest, @@ -466,12 +544,35 @@ async def post_create_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for create_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_artifact` interceptor runs + before the `post_create_artifact_with_metadata` interceptor. """ return response + async def post_create_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_artifact_with_metadata` + interceptor in new development instead of the `post_create_artifact` interceptor. + When both interceptors are used, this `post_create_artifact_with_metadata` interceptor runs after the + `post_create_artifact` interceptor. The (possibly modified) response returned by + `post_create_artifact` will be passed to + `post_create_artifact_with_metadata`. + """ + return response, metadata + async def pre_create_context( self, request: metadata_service.CreateContextRequest, @@ -491,12 +592,35 @@ async def post_create_context( ) -> gca_context.Context: """Post-rpc interceptor for create_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_context` interceptor runs + before the `post_create_context_with_metadata` interceptor. """ return response + async def post_create_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_context_with_metadata` + interceptor in new development instead of the `post_create_context` interceptor. + When both interceptors are used, this `post_create_context_with_metadata` interceptor runs after the + `post_create_context` interceptor. The (possibly modified) response returned by + `post_create_context` will be passed to + `post_create_context_with_metadata`. + """ + return response, metadata + async def pre_create_execution( self, request: metadata_service.CreateExecutionRequest, @@ -516,12 +640,35 @@ async def post_create_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for create_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_execution` interceptor runs + before the `post_create_execution_with_metadata` interceptor. """ return response + async def post_create_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_execution_with_metadata` + interceptor in new development instead of the `post_create_execution` interceptor. + When both interceptors are used, this `post_create_execution_with_metadata` interceptor runs after the + `post_create_execution` interceptor. The (possibly modified) response returned by + `post_create_execution` will be passed to + `post_create_execution_with_metadata`. + """ + return response, metadata + async def pre_create_metadata_schema( self, request: metadata_service.CreateMetadataSchemaRequest, @@ -542,12 +689,37 @@ async def post_create_metadata_schema( ) -> gca_metadata_schema.MetadataSchema: """Post-rpc interceptor for create_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_schema` interceptor runs + before the `post_create_metadata_schema_with_metadata` interceptor. """ return response + async def post_create_metadata_schema_with_metadata( + self, + response: gca_metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_schema_with_metadata` + interceptor in new development instead of the `post_create_metadata_schema` interceptor. + When both interceptors are used, this `post_create_metadata_schema_with_metadata` interceptor runs after the + `post_create_metadata_schema` interceptor. The (possibly modified) response returned by + `post_create_metadata_schema` will be passed to + `post_create_metadata_schema_with_metadata`. + """ + return response, metadata + async def pre_create_metadata_store( self, request: metadata_service.CreateMetadataStoreRequest, @@ -568,12 +740,35 @@ async def post_create_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_store` interceptor runs + before the `post_create_metadata_store_with_metadata` interceptor. """ return response + async def post_create_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_store_with_metadata` + interceptor in new development instead of the `post_create_metadata_store` interceptor. + When both interceptors are used, this `post_create_metadata_store_with_metadata` interceptor runs after the + `post_create_metadata_store` interceptor. The (possibly modified) response returned by + `post_create_metadata_store` will be passed to + `post_create_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_delete_artifact( self, request: metadata_service.DeleteArtifactRequest, @@ -593,12 +788,35 @@ async def post_delete_artifact( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_artifact` interceptor runs + before the `post_delete_artifact_with_metadata` interceptor. """ return response + async def post_delete_artifact_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_artifact_with_metadata` + interceptor in new development instead of the `post_delete_artifact` interceptor. + When both interceptors are used, this `post_delete_artifact_with_metadata` interceptor runs after the + `post_delete_artifact` interceptor. The (possibly modified) response returned by + `post_delete_artifact` will be passed to + `post_delete_artifact_with_metadata`. + """ + return response, metadata + async def pre_delete_context( self, request: metadata_service.DeleteContextRequest, @@ -618,12 +836,35 @@ async def post_delete_context( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_context` interceptor runs + before the `post_delete_context_with_metadata` interceptor. """ return response + async def post_delete_context_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_context_with_metadata` + interceptor in new development instead of the `post_delete_context` interceptor. + When both interceptors are used, this `post_delete_context_with_metadata` interceptor runs after the + `post_delete_context` interceptor. The (possibly modified) response returned by + `post_delete_context` will be passed to + `post_delete_context_with_metadata`. + """ + return response, metadata + async def pre_delete_execution( self, request: metadata_service.DeleteExecutionRequest, @@ -643,12 +884,35 @@ async def post_delete_execution( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_execution` interceptor runs + before the `post_delete_execution_with_metadata` interceptor. """ return response + async def post_delete_execution_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_execution_with_metadata` + interceptor in new development instead of the `post_delete_execution` interceptor. + When both interceptors are used, this `post_delete_execution_with_metadata` interceptor runs after the + `post_delete_execution` interceptor. The (possibly modified) response returned by + `post_delete_execution` will be passed to + `post_delete_execution_with_metadata`. + """ + return response, metadata + async def pre_delete_metadata_store( self, request: metadata_service.DeleteMetadataStoreRequest, @@ -669,12 +933,35 @@ async def post_delete_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_metadata_store` interceptor runs + before the `post_delete_metadata_store_with_metadata` interceptor. """ return response + async def post_delete_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_metadata_store_with_metadata` + interceptor in new development instead of the `post_delete_metadata_store` interceptor. + When both interceptors are used, this `post_delete_metadata_store_with_metadata` interceptor runs after the + `post_delete_metadata_store` interceptor. The (possibly modified) response returned by + `post_delete_metadata_store` will be passed to + `post_delete_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_get_artifact( self, request: metadata_service.GetArtifactRequest, @@ -692,12 +979,35 @@ async def pre_get_artifact( async def post_get_artifact(self, response: artifact.Artifact) -> artifact.Artifact: """Post-rpc interceptor for get_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_artifact` interceptor runs + before the `post_get_artifact_with_metadata` interceptor. """ return response + async def post_get_artifact_with_metadata( + self, + response: artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_artifact_with_metadata` + interceptor in new development instead of the `post_get_artifact` interceptor. + When both interceptors are used, this `post_get_artifact_with_metadata` interceptor runs after the + `post_get_artifact` interceptor. The (possibly modified) response returned by + `post_get_artifact` will be passed to + `post_get_artifact_with_metadata`. + """ + return response, metadata + async def pre_get_context( self, request: metadata_service.GetContextRequest, @@ -715,12 +1025,35 @@ async def pre_get_context( async def post_get_context(self, response: context.Context) -> context.Context: """Post-rpc interceptor for get_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_context` interceptor runs + before the `post_get_context_with_metadata` interceptor. """ return response + async def post_get_context_with_metadata( + self, + response: context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_context_with_metadata` + interceptor in new development instead of the `post_get_context` interceptor. + When both interceptors are used, this `post_get_context_with_metadata` interceptor runs after the + `post_get_context` interceptor. The (possibly modified) response returned by + `post_get_context` will be passed to + `post_get_context_with_metadata`. + """ + return response, metadata + async def pre_get_execution( self, request: metadata_service.GetExecutionRequest, @@ -740,12 +1073,35 @@ async def post_get_execution( ) -> execution.Execution: """Post-rpc interceptor for get_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_execution` interceptor runs + before the `post_get_execution_with_metadata` interceptor. """ return response + async def post_get_execution_with_metadata( + self, + response: execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_execution_with_metadata` + interceptor in new development instead of the `post_get_execution` interceptor. + When both interceptors are used, this `post_get_execution_with_metadata` interceptor runs after the + `post_get_execution` interceptor. The (possibly modified) response returned by + `post_get_execution` will be passed to + `post_get_execution_with_metadata`. + """ + return response, metadata + async def pre_get_metadata_schema( self, request: metadata_service.GetMetadataSchemaRequest, @@ -766,12 +1122,35 @@ async def post_get_metadata_schema( ) -> metadata_schema.MetadataSchema: """Post-rpc interceptor for get_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_schema` interceptor runs + before the `post_get_metadata_schema_with_metadata` interceptor. """ return response + async def post_get_metadata_schema_with_metadata( + self, + response: metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_schema_with_metadata` + interceptor in new development instead of the `post_get_metadata_schema` interceptor. + When both interceptors are used, this `post_get_metadata_schema_with_metadata` interceptor runs after the + `post_get_metadata_schema` interceptor. The (possibly modified) response returned by + `post_get_metadata_schema` will be passed to + `post_get_metadata_schema_with_metadata`. + """ + return response, metadata + async def pre_get_metadata_store( self, request: metadata_service.GetMetadataStoreRequest, @@ -792,12 +1171,35 @@ async def post_get_metadata_store( ) -> metadata_store.MetadataStore: """Post-rpc interceptor for get_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_store` interceptor runs + before the `post_get_metadata_store_with_metadata` interceptor. """ return response + async def post_get_metadata_store_with_metadata( + self, + response: metadata_store.MetadataStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_store.MetadataStore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_store_with_metadata` + interceptor in new development instead of the `post_get_metadata_store` interceptor. + When both interceptors are used, this `post_get_metadata_store_with_metadata` interceptor runs after the + `post_get_metadata_store` interceptor. The (possibly modified) response returned by + `post_get_metadata_store` will be passed to + `post_get_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_list_artifacts( self, request: metadata_service.ListArtifactsRequest, @@ -817,12 +1219,37 @@ async def post_list_artifacts( ) -> metadata_service.ListArtifactsResponse: """Post-rpc interceptor for list_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_artifacts` interceptor runs + before the `post_list_artifacts_with_metadata` interceptor. """ return response + async def post_list_artifacts_with_metadata( + self, + response: metadata_service.ListArtifactsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListArtifactsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_artifacts_with_metadata` + interceptor in new development instead of the `post_list_artifacts` interceptor. + When both interceptors are used, this `post_list_artifacts_with_metadata` interceptor runs after the + `post_list_artifacts` interceptor. The (possibly modified) response returned by + `post_list_artifacts` will be passed to + `post_list_artifacts_with_metadata`. + """ + return response, metadata + async def pre_list_contexts( self, request: metadata_service.ListContextsRequest, @@ -842,12 +1269,37 @@ async def post_list_contexts( ) -> metadata_service.ListContextsResponse: """Post-rpc interceptor for list_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_contexts` interceptor runs + before the `post_list_contexts_with_metadata` interceptor. """ return response + async def post_list_contexts_with_metadata( + self, + response: metadata_service.ListContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListContextsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_contexts_with_metadata` + interceptor in new development instead of the `post_list_contexts` interceptor. + When both interceptors are used, this `post_list_contexts_with_metadata` interceptor runs after the + `post_list_contexts` interceptor. The (possibly modified) response returned by + `post_list_contexts` will be passed to + `post_list_contexts_with_metadata`. + """ + return response, metadata + async def pre_list_executions( self, request: metadata_service.ListExecutionsRequest, @@ -867,12 +1319,37 @@ async def post_list_executions( ) -> metadata_service.ListExecutionsResponse: """Post-rpc interceptor for list_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_executions` interceptor runs + before the `post_list_executions_with_metadata` interceptor. """ return response + async def post_list_executions_with_metadata( + self, + response: metadata_service.ListExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListExecutionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_executions_with_metadata` + interceptor in new development instead of the `post_list_executions` interceptor. + When both interceptors are used, this `post_list_executions_with_metadata` interceptor runs after the + `post_list_executions` interceptor. The (possibly modified) response returned by + `post_list_executions` will be passed to + `post_list_executions_with_metadata`. + """ + return response, metadata + async def pre_list_metadata_schemas( self, request: metadata_service.ListMetadataSchemasRequest, @@ -893,12 +1370,38 @@ async def post_list_metadata_schemas( ) -> metadata_service.ListMetadataSchemasResponse: """Post-rpc interceptor for list_metadata_schemas - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_schemas_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_schemas` interceptor runs + before the `post_list_metadata_schemas_with_metadata` interceptor. """ return response + async def post_list_metadata_schemas_with_metadata( + self, + response: metadata_service.ListMetadataSchemasResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataSchemasResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_schemas + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_schemas_with_metadata` + interceptor in new development instead of the `post_list_metadata_schemas` interceptor. + When both interceptors are used, this `post_list_metadata_schemas_with_metadata` interceptor runs after the + `post_list_metadata_schemas` interceptor. The (possibly modified) response returned by + `post_list_metadata_schemas` will be passed to + `post_list_metadata_schemas_with_metadata`. + """ + return response, metadata + async def pre_list_metadata_stores( self, request: metadata_service.ListMetadataStoresRequest, @@ -919,12 +1422,38 @@ async def post_list_metadata_stores( ) -> metadata_service.ListMetadataStoresResponse: """Post-rpc interceptor for list_metadata_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_stores` interceptor runs + before the `post_list_metadata_stores_with_metadata` interceptor. """ return response + async def post_list_metadata_stores_with_metadata( + self, + response: metadata_service.ListMetadataStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_stores_with_metadata` + interceptor in new development instead of the `post_list_metadata_stores` interceptor. + When both interceptors are used, this `post_list_metadata_stores_with_metadata` interceptor runs after the + `post_list_metadata_stores` interceptor. The (possibly modified) response returned by + `post_list_metadata_stores` will be passed to + `post_list_metadata_stores_with_metadata`. + """ + return response, metadata + async def pre_purge_artifacts( self, request: metadata_service.PurgeArtifactsRequest, @@ -944,12 +1473,35 @@ async def post_purge_artifacts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_artifacts` interceptor runs + before the `post_purge_artifacts_with_metadata` interceptor. """ return response + async def post_purge_artifacts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_artifacts_with_metadata` + interceptor in new development instead of the `post_purge_artifacts` interceptor. + When both interceptors are used, this `post_purge_artifacts_with_metadata` interceptor runs after the + `post_purge_artifacts` interceptor. The (possibly modified) response returned by + `post_purge_artifacts` will be passed to + `post_purge_artifacts_with_metadata`. + """ + return response, metadata + async def pre_purge_contexts( self, request: metadata_service.PurgeContextsRequest, @@ -969,12 +1521,35 @@ async def post_purge_contexts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_contexts` interceptor runs + before the `post_purge_contexts_with_metadata` interceptor. """ return response + async def post_purge_contexts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_contexts_with_metadata` + interceptor in new development instead of the `post_purge_contexts` interceptor. + When both interceptors are used, this `post_purge_contexts_with_metadata` interceptor runs after the + `post_purge_contexts` interceptor. The (possibly modified) response returned by + `post_purge_contexts` will be passed to + `post_purge_contexts_with_metadata`. + """ + return response, metadata + async def pre_purge_executions( self, request: metadata_service.PurgeExecutionsRequest, @@ -994,12 +1569,35 @@ async def post_purge_executions( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_executions` interceptor runs + before the `post_purge_executions_with_metadata` interceptor. """ return response + async def post_purge_executions_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_executions_with_metadata` + interceptor in new development instead of the `post_purge_executions` interceptor. + When both interceptors are used, this `post_purge_executions_with_metadata` interceptor runs after the + `post_purge_executions` interceptor. The (possibly modified) response returned by + `post_purge_executions` will be passed to + `post_purge_executions_with_metadata`. + """ + return response, metadata + async def pre_query_artifact_lineage_subgraph( self, request: metadata_service.QueryArtifactLineageSubgraphRequest, @@ -1020,12 +1618,37 @@ async def post_query_artifact_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_artifact_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_artifact_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_artifact_lineage_subgraph` interceptor runs + before the `post_query_artifact_lineage_subgraph_with_metadata` interceptor. """ return response + async def post_query_artifact_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_artifact_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_artifact_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_artifact_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_artifact_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_artifact_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_artifact_lineage_subgraph` will be passed to + `post_query_artifact_lineage_subgraph_with_metadata`. + """ + return response, metadata + async def pre_query_context_lineage_subgraph( self, request: metadata_service.QueryContextLineageSubgraphRequest, @@ -1046,12 +1669,37 @@ async def post_query_context_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_context_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_context_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_context_lineage_subgraph` interceptor runs + before the `post_query_context_lineage_subgraph_with_metadata` interceptor. """ return response + async def post_query_context_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_context_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_context_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_context_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_context_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_context_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_context_lineage_subgraph` will be passed to + `post_query_context_lineage_subgraph_with_metadata`. + """ + return response, metadata + async def pre_query_execution_inputs_and_outputs( self, request: metadata_service.QueryExecutionInputsAndOutputsRequest, @@ -1072,12 +1720,37 @@ async def post_query_execution_inputs_and_outputs( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_execution_inputs_and_outputs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_execution_inputs_and_outputs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_execution_inputs_and_outputs` interceptor runs + before the `post_query_execution_inputs_and_outputs_with_metadata` interceptor. """ return response + async def post_query_execution_inputs_and_outputs_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_execution_inputs_and_outputs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_execution_inputs_and_outputs_with_metadata` + interceptor in new development instead of the `post_query_execution_inputs_and_outputs` interceptor. + When both interceptors are used, this `post_query_execution_inputs_and_outputs_with_metadata` interceptor runs after the + `post_query_execution_inputs_and_outputs` interceptor. The (possibly modified) response returned by + `post_query_execution_inputs_and_outputs` will be passed to + `post_query_execution_inputs_and_outputs_with_metadata`. + """ + return response, metadata + async def pre_remove_context_children( self, request: metadata_service.RemoveContextChildrenRequest, @@ -1098,12 +1771,38 @@ async def post_remove_context_children( ) -> metadata_service.RemoveContextChildrenResponse: """Post-rpc interceptor for remove_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_remove_context_children` interceptor runs + before the `post_remove_context_children_with_metadata` interceptor. """ return response + async def post_remove_context_children_with_metadata( + self, + response: metadata_service.RemoveContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.RemoveContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for remove_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_remove_context_children_with_metadata` + interceptor in new development instead of the `post_remove_context_children` interceptor. + When both interceptors are used, this `post_remove_context_children_with_metadata` interceptor runs after the + `post_remove_context_children` interceptor. The (possibly modified) response returned by + `post_remove_context_children` will be passed to + `post_remove_context_children_with_metadata`. + """ + return response, metadata + async def pre_update_artifact( self, request: metadata_service.UpdateArtifactRequest, @@ -1123,12 +1822,35 @@ async def post_update_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for update_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_artifact` interceptor runs + before the `post_update_artifact_with_metadata` interceptor. """ return response + async def post_update_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_artifact_with_metadata` + interceptor in new development instead of the `post_update_artifact` interceptor. + When both interceptors are used, this `post_update_artifact_with_metadata` interceptor runs after the + `post_update_artifact` interceptor. The (possibly modified) response returned by + `post_update_artifact` will be passed to + `post_update_artifact_with_metadata`. + """ + return response, metadata + async def pre_update_context( self, request: metadata_service.UpdateContextRequest, @@ -1148,12 +1870,35 @@ async def post_update_context( ) -> gca_context.Context: """Post-rpc interceptor for update_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_context` interceptor runs + before the `post_update_context_with_metadata` interceptor. """ return response + async def post_update_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_context_with_metadata` + interceptor in new development instead of the `post_update_context` interceptor. + When both interceptors are used, this `post_update_context_with_metadata` interceptor runs after the + `post_update_context` interceptor. The (possibly modified) response returned by + `post_update_context` will be passed to + `post_update_context_with_metadata`. + """ + return response, metadata + async def pre_update_execution( self, request: metadata_service.UpdateExecutionRequest, @@ -1173,12 +1918,35 @@ async def post_update_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for update_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_execution` interceptor runs + before the `post_update_execution_with_metadata` interceptor. """ return response + async def post_update_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_execution_with_metadata` + interceptor in new development instead of the `post_update_execution` interceptor. + When both interceptors are used, this `post_update_execution_with_metadata` interceptor runs after the + `post_update_execution` interceptor. The (possibly modified) response returned by + `post_update_execution` will be passed to + `post_update_execution_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1854,6 +2622,13 @@ async def __call__( resp = await self._interceptor.post_add_context_artifacts_and_executions( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_add_context_artifacts_and_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2015,6 +2790,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2176,6 +2955,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_execution_events(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_execution_events_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2337,6 +3120,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2496,6 +3283,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2655,6 +3446,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2812,6 +3607,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2974,6 +3773,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3130,6 +3933,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3286,6 +4093,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3442,6 +4253,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3596,6 +4411,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3748,6 +4567,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3900,6 +4723,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4053,6 +4880,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4204,6 +5035,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4360,6 +5195,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4515,6 +5354,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4672,6 +5515,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4829,6 +5676,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4984,6 +5835,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_metadata_schemas(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_metadata_schemas_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5139,6 +5994,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_metadata_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_metadata_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5303,6 +6162,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5465,6 +6328,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5627,6 +6494,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5787,6 +6658,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_artifact_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_artifact_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5947,6 +6825,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_context_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_context_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6109,6 +6994,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_execution_inputs_and_outputs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_execution_inputs_and_outputs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6270,6 +7162,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_remove_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_remove_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6431,6 +7330,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6590,6 +7493,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6749,6 +7656,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/migration_service/client.py b/google/cloud/aiplatform_v1/services/migration_service/client.py index ce85c8e49f..ad7b08c916 100644 --- a/google/cloud/aiplatform_v1/services/migration_service/client.py +++ b/google/cloud/aiplatform_v1/services/migration_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -240,40 +242,40 @@ def parse_annotated_dataset_path(path: str) -> Dict[str, str]: @staticmethod def dataset_path( project: str, - location: str, dataset: str, ) -> str: """Returns a fully-qualified dataset string.""" - return "projects/{project}/locations/{location}/datasets/{dataset}".format( + return "projects/{project}/datasets/{dataset}".format( project=project, - location=location, dataset=dataset, ) @staticmethod def parse_dataset_path(path: str) -> Dict[str, str]: """Parses a dataset path into its component segments.""" - m = re.match( - r"^projects/(?P.+?)/locations/(?P.+?)/datasets/(?P.+?)$", - path, - ) + m = re.match(r"^projects/(?P.+?)/datasets/(?P.+?)$", path) return m.groupdict() if m else {} @staticmethod def dataset_path( project: str, + location: str, dataset: str, ) -> str: """Returns a fully-qualified dataset string.""" - return "projects/{project}/datasets/{dataset}".format( + return "projects/{project}/locations/{location}/datasets/{dataset}".format( project=project, + location=location, dataset=dataset, ) @staticmethod def parse_dataset_path(path: str) -> Dict[str, str]: """Parses a dataset path into its component segments.""" - m = re.match(r"^projects/(?P.+?)/datasets/(?P.+?)$", path) + m = re.match( + r"^projects/(?P.+?)/locations/(?P.+?)/datasets/(?P.+?)$", + path, + ) return m.groupdict() if m else {} @staticmethod @@ -633,6 +635,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1189,16 +1218,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1244,16 +1277,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1416,16 +1453,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1537,16 +1578,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1659,16 +1704,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1719,16 +1768,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1774,16 +1827,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1829,16 +1886,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/migration_service/transports/rest.py b/google/cloud/aiplatform_v1/services/migration_service/transports/rest.py index 02641f3d6f..0e78a2469f 100644 --- a/google/cloud/aiplatform_v1/services/migration_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/migration_service/transports/rest.py @@ -121,12 +121,35 @@ def post_batch_migrate_resources( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_migrate_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_migrate_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_batch_migrate_resources` interceptor runs + before the `post_batch_migrate_resources_with_metadata` interceptor. """ return response + def post_batch_migrate_resources_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_migrate_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_batch_migrate_resources_with_metadata` + interceptor in new development instead of the `post_batch_migrate_resources` interceptor. + When both interceptors are used, this `post_batch_migrate_resources_with_metadata` interceptor runs after the + `post_batch_migrate_resources` interceptor. The (possibly modified) response returned by + `post_batch_migrate_resources` will be passed to + `post_batch_migrate_resources_with_metadata`. + """ + return response, metadata + def pre_search_migratable_resources( self, request: migration_service.SearchMigratableResourcesRequest, @@ -147,12 +170,38 @@ def post_search_migratable_resources( ) -> migration_service.SearchMigratableResourcesResponse: """Post-rpc interceptor for search_migratable_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_migratable_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_search_migratable_resources` interceptor runs + before the `post_search_migratable_resources_with_metadata` interceptor. """ return response + def post_search_migratable_resources_with_metadata( + self, + response: migration_service.SearchMigratableResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + migration_service.SearchMigratableResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_migratable_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_search_migratable_resources_with_metadata` + interceptor in new development instead of the `post_search_migratable_resources` interceptor. + When both interceptors are used, this `post_search_migratable_resources_with_metadata` interceptor runs after the + `post_search_migratable_resources` interceptor. The (possibly modified) response returned by + `post_search_migratable_resources` will be passed to + `post_search_migratable_resources_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2530,6 +2579,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_migrate_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_migrate_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2686,6 +2739,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_migratable_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_migratable_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/migration_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/migration_service/transports/rest_asyncio.py index b72ec31ec7..fb1adc6efb 100644 --- a/google/cloud/aiplatform_v1/services/migration_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/migration_service/transports/rest_asyncio.py @@ -139,12 +139,35 @@ async def post_batch_migrate_resources( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_migrate_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_migrate_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_batch_migrate_resources` interceptor runs + before the `post_batch_migrate_resources_with_metadata` interceptor. """ return response + async def post_batch_migrate_resources_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_migrate_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_batch_migrate_resources_with_metadata` + interceptor in new development instead of the `post_batch_migrate_resources` interceptor. + When both interceptors are used, this `post_batch_migrate_resources_with_metadata` interceptor runs after the + `post_batch_migrate_resources` interceptor. The (possibly modified) response returned by + `post_batch_migrate_resources` will be passed to + `post_batch_migrate_resources_with_metadata`. + """ + return response, metadata + async def pre_search_migratable_resources( self, request: migration_service.SearchMigratableResourcesRequest, @@ -165,12 +188,38 @@ async def post_search_migratable_resources( ) -> migration_service.SearchMigratableResourcesResponse: """Post-rpc interceptor for search_migratable_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_migratable_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_search_migratable_resources` interceptor runs + before the `post_search_migratable_resources_with_metadata` interceptor. """ return response + async def post_search_migratable_resources_with_metadata( + self, + response: migration_service.SearchMigratableResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + migration_service.SearchMigratableResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_migratable_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_search_migratable_resources_with_metadata` + interceptor in new development instead of the `post_search_migratable_resources` interceptor. + When both interceptors are used, this `post_search_migratable_resources_with_metadata` interceptor runs after the + `post_search_migratable_resources` interceptor. The (possibly modified) response returned by + `post_search_migratable_resources` will be passed to + `post_search_migratable_resources_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -690,6 +739,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_migrate_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_migrate_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -850,6 +906,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_migratable_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_migratable_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/model_garden_service/client.py b/google/cloud/aiplatform_v1/services/model_garden_service/client.py index 317dfae228..70d9b62b7c 100644 --- a/google/cloud/aiplatform_v1/services/model_garden_service/client.py +++ b/google/cloud/aiplatform_v1/services/model_garden_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -519,6 +521,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -908,16 +937,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -963,16 +996,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1135,16 +1172,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1256,16 +1297,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1378,16 +1423,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1438,16 +1487,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1493,16 +1546,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1548,16 +1605,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest.py b/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest.py index 3043dba7ca..74fc60885d 100644 --- a/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest.py @@ -113,12 +113,35 @@ def post_get_publisher_model( ) -> publisher_model.PublisherModel: """Post-rpc interceptor for get_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_get_publisher_model` interceptor runs + before the `post_get_publisher_model_with_metadata` interceptor. """ return response + def post_get_publisher_model_with_metadata( + self, + response: publisher_model.PublisherModel, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[publisher_model.PublisherModel, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_get_publisher_model_with_metadata` + interceptor in new development instead of the `post_get_publisher_model` interceptor. + When both interceptors are used, this `post_get_publisher_model_with_metadata` interceptor runs after the + `post_get_publisher_model` interceptor. The (possibly modified) response returned by + `post_get_publisher_model` will be passed to + `post_get_publisher_model_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -571,6 +594,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest_asyncio.py index 2104caf6f4..c8151fbeae 100644 --- a/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/model_garden_service/transports/rest_asyncio.py @@ -130,12 +130,35 @@ async def post_get_publisher_model( ) -> publisher_model.PublisherModel: """Post-rpc interceptor for get_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_get_publisher_model` interceptor runs + before the `post_get_publisher_model_with_metadata` interceptor. """ return response + async def post_get_publisher_model_with_metadata( + self, + response: publisher_model.PublisherModel, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[publisher_model.PublisherModel, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_get_publisher_model_with_metadata` + interceptor in new development instead of the `post_get_publisher_model` interceptor. + When both interceptors are used, this `post_get_publisher_model_with_metadata` interceptor runs after the + `post_get_publisher_model` interceptor. The (possibly modified) response returned by + `post_get_publisher_model` will be passed to + `post_get_publisher_model_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -637,6 +660,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/model_service/async_client.py b/google/cloud/aiplatform_v1/services/model_service/async_client.py index 44f08473ef..42805f025d 100644 --- a/google/cloud/aiplatform_v1/services/model_service/async_client.py +++ b/google/cloud/aiplatform_v1/services/model_service/async_client.py @@ -820,6 +820,138 @@ async def sample_list_model_versions(): # Done; return the response. return response + async def list_model_version_checkpoints( + self, + request: Optional[ + Union[model_service.ListModelVersionCheckpointsRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> pagers.ListModelVersionCheckpointsAsyncPager: + r"""Lists checkpoints of the specified model version. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1 + + async def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1.ModelServiceAsyncClient() + + # Initialize request argument(s) + request = aiplatform_v1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + async for response in page_result: + print(response) + + Args: + request (Optional[Union[google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest, dict]]): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints]. + name (:class:`str`): + Required. The name of the model version to list + checkpoints for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest + version will be used. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.cloud.aiplatform_v1.services.model_service.pagers.ListModelVersionCheckpointsAsyncPager: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + + Iterating over this object will yield results and + resolve additional pages automatically. + + """ + # Create or coerce a protobuf request object. + # - Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, model_service.ListModelVersionCheckpointsRequest): + request = model_service.ListModelVersionCheckpointsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._client._transport._wrapped_methods[ + self._client._transport.list_model_version_checkpoints + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # This method is paged; wrap the response in a pager, which provides + # an `__aiter__` convenience method. + response = pagers.ListModelVersionCheckpointsAsyncPager( + method=rpc, + request=request, + response=response, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + async def update_model( self, request: Optional[Union[model_service.UpdateModelRequest, dict]] = None, diff --git a/google/cloud/aiplatform_v1/services/model_service/client.py b/google/cloud/aiplatform_v1/services/model_service/client.py index e77f13eb48..24aa590355 100644 --- a/google/cloud/aiplatform_v1/services/model_service/client.py +++ b/google/cloud/aiplatform_v1/services/model_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -631,6 +633,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1337,6 +1366,137 @@ def sample_list_model_versions(): # Done; return the response. return response + def list_model_version_checkpoints( + self, + request: Optional[ + Union[model_service.ListModelVersionCheckpointsRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> pagers.ListModelVersionCheckpointsPager: + r"""Lists checkpoints of the specified model version. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1 + + def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1.ModelServiceClient() + + # Initialize request argument(s) + request = aiplatform_v1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + for response in page_result: + print(response) + + Args: + request (Union[google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest, dict]): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints]. + name (str): + Required. The name of the model version to list + checkpoints for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest + version will be used. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.cloud.aiplatform_v1.services.model_service.pagers.ListModelVersionCheckpointsPager: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + + Iterating over this object will yield results and + resolve additional pages automatically. + + """ + # Create or coerce a protobuf request object. + # - Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, model_service.ListModelVersionCheckpointsRequest): + request = model_service.ListModelVersionCheckpointsRequest(request) + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[ + self._transport.list_model_version_checkpoints + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # This method is paged; wrap the response in a pager, which provides + # an `__iter__` convenience method. + response = pagers.ListModelVersionCheckpointsPager( + method=rpc, + request=request, + response=response, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + def update_model( self, request: Optional[Union[model_service.UpdateModelRequest, dict]] = None, @@ -3175,16 +3335,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3230,16 +3394,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3402,16 +3570,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3523,16 +3695,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3645,16 +3821,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3705,16 +3885,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3760,16 +3944,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3815,16 +4003,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/model_service/pagers.py b/google/cloud/aiplatform_v1/services/model_service/pagers.py index 03d454f889..3b7e0fd312 100644 --- a/google/cloud/aiplatform_v1/services/model_service/pagers.py +++ b/google/cloud/aiplatform_v1/services/model_service/pagers.py @@ -355,6 +355,166 @@ def __repr__(self) -> str: return "{0}<{1!r}>".format(self.__class__.__name__, self._response) +class ListModelVersionCheckpointsPager: + """A pager for iterating through ``list_model_version_checkpoints`` requests. + + This class thinly wraps an initial + :class:`google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse` object, and + provides an ``__iter__`` method to iterate through its + ``checkpoints`` field. + + If there are more pages, the ``__iter__`` method will make additional + ``ListModelVersionCheckpoints`` requests and continue to iterate + through the ``checkpoints`` field on the + corresponding responses. + + All the usual :class:`google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse` + attributes are available on the pager. If multiple requests are made, only + the most recent response is retained, and thus used for attribute lookup. + """ + + def __init__( + self, + method: Callable[..., model_service.ListModelVersionCheckpointsResponse], + request: model_service.ListModelVersionCheckpointsRequest, + response: model_service.ListModelVersionCheckpointsResponse, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = () + ): + """Instantiate the pager. + + Args: + method (Callable): The method that was originally called, and + which instantiated this pager. + request (google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest): + The initial request object. + response (google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse): + The initial response object. + retry (google.api_core.retry.Retry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + """ + self._method = method + self._request = model_service.ListModelVersionCheckpointsRequest(request) + self._response = response + self._retry = retry + self._timeout = timeout + self._metadata = metadata + + def __getattr__(self, name: str) -> Any: + return getattr(self._response, name) + + @property + def pages(self) -> Iterator[model_service.ListModelVersionCheckpointsResponse]: + yield self._response + while self._response.next_page_token: + self._request.page_token = self._response.next_page_token + self._response = self._method( + self._request, + retry=self._retry, + timeout=self._timeout, + metadata=self._metadata, + ) + yield self._response + + def __iter__(self) -> Iterator[model_service.ModelVersionCheckpoint]: + for page in self.pages: + yield from page.checkpoints + + def __repr__(self) -> str: + return "{0}<{1!r}>".format(self.__class__.__name__, self._response) + + +class ListModelVersionCheckpointsAsyncPager: + """A pager for iterating through ``list_model_version_checkpoints`` requests. + + This class thinly wraps an initial + :class:`google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse` object, and + provides an ``__aiter__`` method to iterate through its + ``checkpoints`` field. + + If there are more pages, the ``__aiter__`` method will make additional + ``ListModelVersionCheckpoints`` requests and continue to iterate + through the ``checkpoints`` field on the + corresponding responses. + + All the usual :class:`google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse` + attributes are available on the pager. If multiple requests are made, only + the most recent response is retained, and thus used for attribute lookup. + """ + + def __init__( + self, + method: Callable[ + ..., Awaitable[model_service.ListModelVersionCheckpointsResponse] + ], + request: model_service.ListModelVersionCheckpointsRequest, + response: model_service.ListModelVersionCheckpointsResponse, + *, + retry: OptionalAsyncRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = () + ): + """Instantiates the pager. + + Args: + method (Callable): The method that was originally called, and + which instantiated this pager. + request (google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest): + The initial request object. + response (google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsResponse): + The initial response object. + retry (google.api_core.retry.AsyncRetry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + """ + self._method = method + self._request = model_service.ListModelVersionCheckpointsRequest(request) + self._response = response + self._retry = retry + self._timeout = timeout + self._metadata = metadata + + def __getattr__(self, name: str) -> Any: + return getattr(self._response, name) + + @property + async def pages( + self, + ) -> AsyncIterator[model_service.ListModelVersionCheckpointsResponse]: + yield self._response + while self._response.next_page_token: + self._request.page_token = self._response.next_page_token + self._response = await self._method( + self._request, + retry=self._retry, + timeout=self._timeout, + metadata=self._metadata, + ) + yield self._response + + def __aiter__(self) -> AsyncIterator[model_service.ModelVersionCheckpoint]: + async def async_generator(): + async for page in self.pages: + for response in page.checkpoints: + yield response + + return async_generator() + + def __repr__(self) -> str: + return "{0}<{1!r}>".format(self.__class__.__name__, self._response) + + class ListModelEvaluationsPager: """A pager for iterating through ``list_model_evaluations`` requests. diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/base.py b/google/cloud/aiplatform_v1/services/model_service/transports/base.py index fac38705d6..8289cffa89 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/base.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/base.py @@ -159,6 +159,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: gapic_v1.method.wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: gapic_v1.method.wrap_method( self.update_model, default_timeout=None, @@ -336,6 +341,18 @@ def list_model_versions( ]: raise NotImplementedError() + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + Union[ + model_service.ListModelVersionCheckpointsResponse, + Awaitable[model_service.ListModelVersionCheckpointsResponse], + ], + ]: + raise NotImplementedError() + @property def update_model( self, diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/grpc.py b/google/cloud/aiplatform_v1/services/model_service/transports/grpc.py index 9a31d4c9c9..8c6f3f9243 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/grpc.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/grpc.py @@ -453,6 +453,37 @@ def list_model_versions( ) return self._stubs["list_model_versions"] + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + r"""Return a callable for the list model version checkpoints method over gRPC. + + Lists checkpoints of the specified model version. + + Returns: + Callable[[~.ListModelVersionCheckpointsRequest], + ~.ListModelVersionCheckpointsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_model_version_checkpoints" not in self._stubs: + self._stubs[ + "list_model_version_checkpoints" + ] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1.ModelService/ListModelVersionCheckpoints", + request_serializer=model_service.ListModelVersionCheckpointsRequest.serialize, + response_deserializer=model_service.ListModelVersionCheckpointsResponse.deserialize, + ) + return self._stubs["list_model_version_checkpoints"] + @property def update_model( self, diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/grpc_asyncio.py b/google/cloud/aiplatform_v1/services/model_service/transports/grpc_asyncio.py index 51cbbcc959..c65a2a5518 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/grpc_asyncio.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/grpc_asyncio.py @@ -466,6 +466,37 @@ def list_model_versions( ) return self._stubs["list_model_versions"] + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + Awaitable[model_service.ListModelVersionCheckpointsResponse], + ]: + r"""Return a callable for the list model version checkpoints method over gRPC. + + Lists checkpoints of the specified model version. + + Returns: + Callable[[~.ListModelVersionCheckpointsRequest], + Awaitable[~.ListModelVersionCheckpointsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_model_version_checkpoints" not in self._stubs: + self._stubs[ + "list_model_version_checkpoints" + ] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1.ModelService/ListModelVersionCheckpoints", + request_serializer=model_service.ListModelVersionCheckpointsRequest.serialize, + response_deserializer=model_service.ListModelVersionCheckpointsResponse.deserialize, + ) + return self._stubs["list_model_version_checkpoints"] + @property def update_model( self, @@ -922,6 +953,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: self._wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: self._wrap_method( self.update_model, default_timeout=None, diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/rest.py b/google/cloud/aiplatform_v1/services/model_service/transports/rest.py index 0135158ead..13b04394a0 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/rest.py @@ -188,6 +188,14 @@ def post_list_models(self, response): logging.log(f"Received response: {response}") return response + def pre_list_model_version_checkpoints(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_list_model_version_checkpoints(self, response): + logging.log(f"Received response: {response}") + return response + def pre_list_model_versions(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -254,12 +262,38 @@ def post_batch_import_evaluated_annotations( ) -> model_service.BatchImportEvaluatedAnnotationsResponse: """Post-rpc interceptor for batch_import_evaluated_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_evaluated_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_evaluated_annotations` interceptor runs + before the `post_batch_import_evaluated_annotations_with_metadata` interceptor. """ return response + def post_batch_import_evaluated_annotations_with_metadata( + self, + response: model_service.BatchImportEvaluatedAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportEvaluatedAnnotationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_evaluated_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_evaluated_annotations_with_metadata` + interceptor in new development instead of the `post_batch_import_evaluated_annotations` interceptor. + When both interceptors are used, this `post_batch_import_evaluated_annotations_with_metadata` interceptor runs after the + `post_batch_import_evaluated_annotations` interceptor. The (possibly modified) response returned by + `post_batch_import_evaluated_annotations` will be passed to + `post_batch_import_evaluated_annotations_with_metadata`. + """ + return response, metadata + def pre_batch_import_model_evaluation_slices( self, request: model_service.BatchImportModelEvaluationSlicesRequest, @@ -280,12 +314,38 @@ def post_batch_import_model_evaluation_slices( ) -> model_service.BatchImportModelEvaluationSlicesResponse: """Post-rpc interceptor for batch_import_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_model_evaluation_slices` interceptor runs + before the `post_batch_import_model_evaluation_slices_with_metadata` interceptor. """ return response + def post_batch_import_model_evaluation_slices_with_metadata( + self, + response: model_service.BatchImportModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_batch_import_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_batch_import_model_evaluation_slices_with_metadata` interceptor runs after the + `post_batch_import_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_batch_import_model_evaluation_slices` will be passed to + `post_batch_import_model_evaluation_slices_with_metadata`. + """ + return response, metadata + def pre_copy_model( self, request: model_service.CopyModelRequest, @@ -303,12 +363,35 @@ def post_copy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for copy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_copy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_copy_model` interceptor runs + before the `post_copy_model_with_metadata` interceptor. """ return response + def post_copy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for copy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_copy_model_with_metadata` + interceptor in new development instead of the `post_copy_model` interceptor. + When both interceptors are used, this `post_copy_model_with_metadata` interceptor runs after the + `post_copy_model` interceptor. The (possibly modified) response returned by + `post_copy_model` will be passed to + `post_copy_model_with_metadata`. + """ + return response, metadata + def pre_delete_model( self, request: model_service.DeleteModelRequest, @@ -328,12 +411,35 @@ def post_delete_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model` interceptor runs + before the `post_delete_model_with_metadata` interceptor. """ return response + def post_delete_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_with_metadata` + interceptor in new development instead of the `post_delete_model` interceptor. + When both interceptors are used, this `post_delete_model_with_metadata` interceptor runs after the + `post_delete_model` interceptor. The (possibly modified) response returned by + `post_delete_model` will be passed to + `post_delete_model_with_metadata`. + """ + return response, metadata + def pre_delete_model_version( self, request: model_service.DeleteModelVersionRequest, @@ -353,12 +459,35 @@ def post_delete_model_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_version` interceptor runs + before the `post_delete_model_version_with_metadata` interceptor. """ return response + def post_delete_model_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_version_with_metadata` + interceptor in new development instead of the `post_delete_model_version` interceptor. + When both interceptors are used, this `post_delete_model_version_with_metadata` interceptor runs after the + `post_delete_model_version` interceptor. The (possibly modified) response returned by + `post_delete_model_version` will be passed to + `post_delete_model_version_with_metadata`. + """ + return response, metadata + def pre_export_model( self, request: model_service.ExportModelRequest, @@ -378,12 +507,35 @@ def post_export_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_export_model` interceptor runs + before the `post_export_model_with_metadata` interceptor. """ return response + def post_export_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_export_model_with_metadata` + interceptor in new development instead of the `post_export_model` interceptor. + When both interceptors are used, this `post_export_model_with_metadata` interceptor runs after the + `post_export_model` interceptor. The (possibly modified) response returned by + `post_export_model` will be passed to + `post_export_model_with_metadata`. + """ + return response, metadata + def pre_get_model( self, request: model_service.GetModelRequest, @@ -399,12 +551,33 @@ def pre_get_model( def post_get_model(self, response: model.Model) -> model.Model: """Post-rpc interceptor for get_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model` interceptor runs + before the `post_get_model_with_metadata` interceptor. """ return response + def post_get_model_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_with_metadata` + interceptor in new development instead of the `post_get_model` interceptor. + When both interceptors are used, this `post_get_model_with_metadata` interceptor runs after the + `post_get_model` interceptor. The (possibly modified) response returned by + `post_get_model` will be passed to + `post_get_model_with_metadata`. + """ + return response, metadata + def pre_get_model_evaluation( self, request: model_service.GetModelEvaluationRequest, @@ -424,12 +597,37 @@ def post_get_model_evaluation( ) -> model_evaluation.ModelEvaluation: """Post-rpc interceptor for get_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation` interceptor runs + before the `post_get_model_evaluation_with_metadata` interceptor. """ return response + def post_get_model_evaluation_with_metadata( + self, + response: model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation` interceptor. + When both interceptors are used, this `post_get_model_evaluation_with_metadata` interceptor runs after the + `post_get_model_evaluation` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation` will be passed to + `post_get_model_evaluation_with_metadata`. + """ + return response, metadata + def pre_get_model_evaluation_slice( self, request: model_service.GetModelEvaluationSliceRequest, @@ -450,12 +648,38 @@ def post_get_model_evaluation_slice( ) -> model_evaluation_slice.ModelEvaluationSlice: """Post-rpc interceptor for get_model_evaluation_slice - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_slice_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation_slice` interceptor runs + before the `post_get_model_evaluation_slice_with_metadata` interceptor. """ return response + def post_get_model_evaluation_slice_with_metadata( + self, + response: model_evaluation_slice.ModelEvaluationSlice, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation_slice.ModelEvaluationSlice, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_evaluation_slice + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_slice_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation_slice` interceptor. + When both interceptors are used, this `post_get_model_evaluation_slice_with_metadata` interceptor runs after the + `post_get_model_evaluation_slice` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation_slice` will be passed to + `post_get_model_evaluation_slice_with_metadata`. + """ + return response, metadata + def pre_import_model_evaluation( self, request: model_service.ImportModelEvaluationRequest, @@ -476,12 +700,37 @@ def post_import_model_evaluation( ) -> gca_model_evaluation.ModelEvaluation: """Post-rpc interceptor for import_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_import_model_evaluation` interceptor runs + before the `post_import_model_evaluation_with_metadata` interceptor. """ return response + def post_import_model_evaluation_with_metadata( + self, + response: gca_model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for import_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_import_model_evaluation_with_metadata` + interceptor in new development instead of the `post_import_model_evaluation` interceptor. + When both interceptors are used, this `post_import_model_evaluation_with_metadata` interceptor runs after the + `post_import_model_evaluation` interceptor. The (possibly modified) response returned by + `post_import_model_evaluation` will be passed to + `post_import_model_evaluation_with_metadata`. + """ + return response, metadata + def pre_list_model_evaluations( self, request: model_service.ListModelEvaluationsRequest, @@ -502,12 +751,38 @@ def post_list_model_evaluations( ) -> model_service.ListModelEvaluationsResponse: """Post-rpc interceptor for list_model_evaluations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluations` interceptor runs + before the `post_list_model_evaluations_with_metadata` interceptor. """ return response + def post_list_model_evaluations_with_metadata( + self, + response: model_service.ListModelEvaluationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluations_with_metadata` + interceptor in new development instead of the `post_list_model_evaluations` interceptor. + When both interceptors are used, this `post_list_model_evaluations_with_metadata` interceptor runs after the + `post_list_model_evaluations` interceptor. The (possibly modified) response returned by + `post_list_model_evaluations` will be passed to + `post_list_model_evaluations_with_metadata`. + """ + return response, metadata + def pre_list_model_evaluation_slices( self, request: model_service.ListModelEvaluationSlicesRequest, @@ -528,12 +803,38 @@ def post_list_model_evaluation_slices( ) -> model_service.ListModelEvaluationSlicesResponse: """Post-rpc interceptor for list_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluation_slices` interceptor runs + before the `post_list_model_evaluation_slices_with_metadata` interceptor. """ return response + def post_list_model_evaluation_slices_with_metadata( + self, + response: model_service.ListModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_list_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_list_model_evaluation_slices_with_metadata` interceptor runs after the + `post_list_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_list_model_evaluation_slices` will be passed to + `post_list_model_evaluation_slices_with_metadata`. + """ + return response, metadata + def pre_list_models( self, request: model_service.ListModelsRequest, @@ -553,12 +854,89 @@ def post_list_models( ) -> model_service.ListModelsResponse: """Post-rpc interceptor for list_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_models` interceptor runs + before the `post_list_models_with_metadata` interceptor. """ return response + def post_list_models_with_metadata( + self, + response: model_service.ListModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_models_with_metadata` + interceptor in new development instead of the `post_list_models` interceptor. + When both interceptors are used, this `post_list_models_with_metadata` interceptor runs after the + `post_list_models` interceptor. The (possibly modified) response returned by + `post_list_models` will be passed to + `post_list_models_with_metadata`. + """ + return response, metadata + + def pre_list_model_version_checkpoints( + self, + request: model_service.ListModelVersionCheckpointsRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to manipulate the request or metadata + before they are sent to the ModelService server. + """ + return request, metadata + + def post_list_model_version_checkpoints( + self, response: model_service.ListModelVersionCheckpointsResponse + ) -> model_service.ListModelVersionCheckpointsResponse: + """Post-rpc interceptor for list_model_version_checkpoints + + DEPRECATED. Please use the `post_list_model_version_checkpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the ModelService server but before + it is returned to user code. This `post_list_model_version_checkpoints` interceptor runs + before the `post_list_model_version_checkpoints_with_metadata` interceptor. + """ + return response + + def post_list_model_version_checkpoints_with_metadata( + self, + response: model_service.ListModelVersionCheckpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_version_checkpoints_with_metadata` + interceptor in new development instead of the `post_list_model_version_checkpoints` interceptor. + When both interceptors are used, this `post_list_model_version_checkpoints_with_metadata` interceptor runs after the + `post_list_model_version_checkpoints` interceptor. The (possibly modified) response returned by + `post_list_model_version_checkpoints` will be passed to + `post_list_model_version_checkpoints_with_metadata`. + """ + return response, metadata + def pre_list_model_versions( self, request: model_service.ListModelVersionsRequest, @@ -578,12 +956,37 @@ def post_list_model_versions( ) -> model_service.ListModelVersionsResponse: """Post-rpc interceptor for list_model_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_versions` interceptor runs + before the `post_list_model_versions_with_metadata` interceptor. """ return response + def post_list_model_versions_with_metadata( + self, + response: model_service.ListModelVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_model_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_versions_with_metadata` + interceptor in new development instead of the `post_list_model_versions` interceptor. + When both interceptors are used, this `post_list_model_versions_with_metadata` interceptor runs after the + `post_list_model_versions` interceptor. The (possibly modified) response returned by + `post_list_model_versions` will be passed to + `post_list_model_versions_with_metadata`. + """ + return response, metadata + def pre_merge_version_aliases( self, request: model_service.MergeVersionAliasesRequest, @@ -602,12 +1005,33 @@ def pre_merge_version_aliases( def post_merge_version_aliases(self, response: model.Model) -> model.Model: """Post-rpc interceptor for merge_version_aliases - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_merge_version_aliases_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_merge_version_aliases` interceptor runs + before the `post_merge_version_aliases_with_metadata` interceptor. """ return response + def post_merge_version_aliases_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for merge_version_aliases + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_merge_version_aliases_with_metadata` + interceptor in new development instead of the `post_merge_version_aliases` interceptor. + When both interceptors are used, this `post_merge_version_aliases_with_metadata` interceptor runs after the + `post_merge_version_aliases` interceptor. The (possibly modified) response returned by + `post_merge_version_aliases` will be passed to + `post_merge_version_aliases_with_metadata`. + """ + return response, metadata + def pre_update_explanation_dataset( self, request: model_service.UpdateExplanationDatasetRequest, @@ -628,12 +1052,35 @@ def post_update_explanation_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_explanation_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_explanation_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_explanation_dataset` interceptor runs + before the `post_update_explanation_dataset_with_metadata` interceptor. """ return response + def post_update_explanation_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_explanation_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_explanation_dataset_with_metadata` + interceptor in new development instead of the `post_update_explanation_dataset` interceptor. + When both interceptors are used, this `post_update_explanation_dataset_with_metadata` interceptor runs after the + `post_update_explanation_dataset` interceptor. The (possibly modified) response returned by + `post_update_explanation_dataset` will be passed to + `post_update_explanation_dataset_with_metadata`. + """ + return response, metadata + def pre_update_model( self, request: model_service.UpdateModelRequest, @@ -651,12 +1098,35 @@ def pre_update_model( def post_update_model(self, response: gca_model.Model) -> gca_model.Model: """Post-rpc interceptor for update_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_model` interceptor runs + before the `post_update_model_with_metadata` interceptor. """ return response + def post_update_model_with_metadata( + self, + response: gca_model.Model, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_model_with_metadata` + interceptor in new development instead of the `post_update_model` interceptor. + When both interceptors are used, this `post_update_model_with_metadata` interceptor runs after the + `post_update_model` interceptor. The (possibly modified) response returned by + `post_update_model` will be passed to + `post_update_model_with_metadata`. + """ + return response, metadata + def pre_upload_model( self, request: model_service.UploadModelRequest, @@ -676,12 +1146,35 @@ def post_upload_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for upload_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_upload_model` interceptor runs + before the `post_upload_model_with_metadata` interceptor. """ return response + def post_upload_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upload_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_upload_model_with_metadata` + interceptor in new development instead of the `post_upload_model` interceptor. + When both interceptors are used, this `post_upload_model_with_metadata` interceptor runs after the + `post_upload_model` interceptor. The (possibly modified) response returned by + `post_upload_model` will be passed to + `post_upload_model_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3061,6 +3554,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_import_evaluated_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_import_evaluated_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3222,6 +3722,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_import_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_import_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3379,6 +3886,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_copy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_copy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3526,6 +4037,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3671,6 +4186,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3826,6 +4345,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3970,6 +4493,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4118,6 +4645,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4270,6 +4801,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_evaluation_slice(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_evaluation_slice_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4426,6 +4961,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4574,6 +5113,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_evaluations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_evaluations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4726,6 +5269,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4878,6 +5425,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4903,6 +5454,167 @@ def __call__( ) return resp + class _ListModelVersionCheckpoints( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints, + ModelServiceRestStub, + ): + def __hash__(self): + return hash("ModelServiceRestTransport.ListModelVersionCheckpoints") + + @staticmethod + def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + ) + return response + + def __call__( + self, + request: model_service.ListModelVersionCheckpointsRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> model_service.ListModelVersionCheckpointsResponse: + r"""Call the list model version + checkpoints method over HTTP. + + Args: + request (~.model_service.ListModelVersionCheckpointsRequest): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints]. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.model_service.ListModelVersionCheckpointsResponse: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + + """ + + http_options = ( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_http_options() + ) + + request, metadata = self._interceptor.pre_list_model_version_checkpoints( + request, metadata + ) + transcoded_request = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_transcoded_request( + http_options, request + ) + + # Jsonify the query params + query_params = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = type(request).to_json(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1.ModelServiceClient.ListModelVersionCheckpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = ( + ModelServiceRestTransport._ListModelVersionCheckpoints._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + ) + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = model_service.ListModelVersionCheckpointsResponse() + pb_resp = model_service.ListModelVersionCheckpointsResponse.pb(resp) + + json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) + + resp = self._interceptor.post_list_model_version_checkpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_model_version_checkpoints_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = ( + model_service.ListModelVersionCheckpointsResponse.to_json( + response + ) + ) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": response.status_code, + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1.ModelServiceClient.list_model_version_checkpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + return resp + class _ListModelVersions( _BaseModelServiceRestTransport._BaseListModelVersions, ModelServiceRestStub ): @@ -5026,6 +5738,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5178,6 +5894,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_merge_version_aliases(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_merge_version_aliases_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5333,6 +6053,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_explanation_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_explanation_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5487,6 +6211,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5642,6 +6370,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_upload_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upload_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5787,6 +6519,17 @@ def list_models( # In C++ this would require a dynamic_cast return self._ListModels(self._session, self._host, self._interceptor) # type: ignore + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._ListModelVersionCheckpoints(self._session, self._host, self._interceptor) # type: ignore + @property def list_model_versions( self, diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/model_service/transports/rest_asyncio.py index a6378edee3..e04167623e 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/rest_asyncio.py @@ -206,6 +206,14 @@ async def post_list_models(self, response): logging.log(f"Received response: {response}") return response + async def pre_list_model_version_checkpoints(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + async def post_list_model_version_checkpoints(self, response): + logging.log(f"Received response: {response}") + return response + async def pre_list_model_versions(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -272,12 +280,38 @@ async def post_batch_import_evaluated_annotations( ) -> model_service.BatchImportEvaluatedAnnotationsResponse: """Post-rpc interceptor for batch_import_evaluated_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_evaluated_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_evaluated_annotations` interceptor runs + before the `post_batch_import_evaluated_annotations_with_metadata` interceptor. """ return response + async def post_batch_import_evaluated_annotations_with_metadata( + self, + response: model_service.BatchImportEvaluatedAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportEvaluatedAnnotationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_evaluated_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_evaluated_annotations_with_metadata` + interceptor in new development instead of the `post_batch_import_evaluated_annotations` interceptor. + When both interceptors are used, this `post_batch_import_evaluated_annotations_with_metadata` interceptor runs after the + `post_batch_import_evaluated_annotations` interceptor. The (possibly modified) response returned by + `post_batch_import_evaluated_annotations` will be passed to + `post_batch_import_evaluated_annotations_with_metadata`. + """ + return response, metadata + async def pre_batch_import_model_evaluation_slices( self, request: model_service.BatchImportModelEvaluationSlicesRequest, @@ -298,12 +332,38 @@ async def post_batch_import_model_evaluation_slices( ) -> model_service.BatchImportModelEvaluationSlicesResponse: """Post-rpc interceptor for batch_import_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_model_evaluation_slices` interceptor runs + before the `post_batch_import_model_evaluation_slices_with_metadata` interceptor. """ return response + async def post_batch_import_model_evaluation_slices_with_metadata( + self, + response: model_service.BatchImportModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_batch_import_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_batch_import_model_evaluation_slices_with_metadata` interceptor runs after the + `post_batch_import_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_batch_import_model_evaluation_slices` will be passed to + `post_batch_import_model_evaluation_slices_with_metadata`. + """ + return response, metadata + async def pre_copy_model( self, request: model_service.CopyModelRequest, @@ -321,12 +381,35 @@ async def post_copy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for copy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_copy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_copy_model` interceptor runs + before the `post_copy_model_with_metadata` interceptor. """ return response + async def post_copy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for copy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_copy_model_with_metadata` + interceptor in new development instead of the `post_copy_model` interceptor. + When both interceptors are used, this `post_copy_model_with_metadata` interceptor runs after the + `post_copy_model` interceptor. The (possibly modified) response returned by + `post_copy_model` will be passed to + `post_copy_model_with_metadata`. + """ + return response, metadata + async def pre_delete_model( self, request: model_service.DeleteModelRequest, @@ -346,12 +429,35 @@ async def post_delete_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model` interceptor runs + before the `post_delete_model_with_metadata` interceptor. """ return response + async def post_delete_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_with_metadata` + interceptor in new development instead of the `post_delete_model` interceptor. + When both interceptors are used, this `post_delete_model_with_metadata` interceptor runs after the + `post_delete_model` interceptor. The (possibly modified) response returned by + `post_delete_model` will be passed to + `post_delete_model_with_metadata`. + """ + return response, metadata + async def pre_delete_model_version( self, request: model_service.DeleteModelVersionRequest, @@ -371,12 +477,35 @@ async def post_delete_model_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_version` interceptor runs + before the `post_delete_model_version_with_metadata` interceptor. """ return response + async def post_delete_model_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_version_with_metadata` + interceptor in new development instead of the `post_delete_model_version` interceptor. + When both interceptors are used, this `post_delete_model_version_with_metadata` interceptor runs after the + `post_delete_model_version` interceptor. The (possibly modified) response returned by + `post_delete_model_version` will be passed to + `post_delete_model_version_with_metadata`. + """ + return response, metadata + async def pre_export_model( self, request: model_service.ExportModelRequest, @@ -396,12 +525,35 @@ async def post_export_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_export_model` interceptor runs + before the `post_export_model_with_metadata` interceptor. """ return response + async def post_export_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_export_model_with_metadata` + interceptor in new development instead of the `post_export_model` interceptor. + When both interceptors are used, this `post_export_model_with_metadata` interceptor runs after the + `post_export_model` interceptor. The (possibly modified) response returned by + `post_export_model` will be passed to + `post_export_model_with_metadata`. + """ + return response, metadata + async def pre_get_model( self, request: model_service.GetModelRequest, @@ -417,12 +569,33 @@ async def pre_get_model( async def post_get_model(self, response: model.Model) -> model.Model: """Post-rpc interceptor for get_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model` interceptor runs + before the `post_get_model_with_metadata` interceptor. """ return response + async def post_get_model_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_with_metadata` + interceptor in new development instead of the `post_get_model` interceptor. + When both interceptors are used, this `post_get_model_with_metadata` interceptor runs after the + `post_get_model` interceptor. The (possibly modified) response returned by + `post_get_model` will be passed to + `post_get_model_with_metadata`. + """ + return response, metadata + async def pre_get_model_evaluation( self, request: model_service.GetModelEvaluationRequest, @@ -442,12 +615,37 @@ async def post_get_model_evaluation( ) -> model_evaluation.ModelEvaluation: """Post-rpc interceptor for get_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation` interceptor runs + before the `post_get_model_evaluation_with_metadata` interceptor. """ return response + async def post_get_model_evaluation_with_metadata( + self, + response: model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation` interceptor. + When both interceptors are used, this `post_get_model_evaluation_with_metadata` interceptor runs after the + `post_get_model_evaluation` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation` will be passed to + `post_get_model_evaluation_with_metadata`. + """ + return response, metadata + async def pre_get_model_evaluation_slice( self, request: model_service.GetModelEvaluationSliceRequest, @@ -468,12 +666,38 @@ async def post_get_model_evaluation_slice( ) -> model_evaluation_slice.ModelEvaluationSlice: """Post-rpc interceptor for get_model_evaluation_slice - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_slice_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation_slice` interceptor runs + before the `post_get_model_evaluation_slice_with_metadata` interceptor. """ return response + async def post_get_model_evaluation_slice_with_metadata( + self, + response: model_evaluation_slice.ModelEvaluationSlice, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation_slice.ModelEvaluationSlice, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_evaluation_slice + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_slice_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation_slice` interceptor. + When both interceptors are used, this `post_get_model_evaluation_slice_with_metadata` interceptor runs after the + `post_get_model_evaluation_slice` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation_slice` will be passed to + `post_get_model_evaluation_slice_with_metadata`. + """ + return response, metadata + async def pre_import_model_evaluation( self, request: model_service.ImportModelEvaluationRequest, @@ -494,12 +718,37 @@ async def post_import_model_evaluation( ) -> gca_model_evaluation.ModelEvaluation: """Post-rpc interceptor for import_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_import_model_evaluation` interceptor runs + before the `post_import_model_evaluation_with_metadata` interceptor. """ return response + async def post_import_model_evaluation_with_metadata( + self, + response: gca_model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for import_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_import_model_evaluation_with_metadata` + interceptor in new development instead of the `post_import_model_evaluation` interceptor. + When both interceptors are used, this `post_import_model_evaluation_with_metadata` interceptor runs after the + `post_import_model_evaluation` interceptor. The (possibly modified) response returned by + `post_import_model_evaluation` will be passed to + `post_import_model_evaluation_with_metadata`. + """ + return response, metadata + async def pre_list_model_evaluations( self, request: model_service.ListModelEvaluationsRequest, @@ -520,12 +769,38 @@ async def post_list_model_evaluations( ) -> model_service.ListModelEvaluationsResponse: """Post-rpc interceptor for list_model_evaluations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluations` interceptor runs + before the `post_list_model_evaluations_with_metadata` interceptor. """ return response + async def post_list_model_evaluations_with_metadata( + self, + response: model_service.ListModelEvaluationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluations_with_metadata` + interceptor in new development instead of the `post_list_model_evaluations` interceptor. + When both interceptors are used, this `post_list_model_evaluations_with_metadata` interceptor runs after the + `post_list_model_evaluations` interceptor. The (possibly modified) response returned by + `post_list_model_evaluations` will be passed to + `post_list_model_evaluations_with_metadata`. + """ + return response, metadata + async def pre_list_model_evaluation_slices( self, request: model_service.ListModelEvaluationSlicesRequest, @@ -546,12 +821,38 @@ async def post_list_model_evaluation_slices( ) -> model_service.ListModelEvaluationSlicesResponse: """Post-rpc interceptor for list_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluation_slices` interceptor runs + before the `post_list_model_evaluation_slices_with_metadata` interceptor. """ return response + async def post_list_model_evaluation_slices_with_metadata( + self, + response: model_service.ListModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_list_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_list_model_evaluation_slices_with_metadata` interceptor runs after the + `post_list_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_list_model_evaluation_slices` will be passed to + `post_list_model_evaluation_slices_with_metadata`. + """ + return response, metadata + async def pre_list_models( self, request: model_service.ListModelsRequest, @@ -571,12 +872,89 @@ async def post_list_models( ) -> model_service.ListModelsResponse: """Post-rpc interceptor for list_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_models` interceptor runs + before the `post_list_models_with_metadata` interceptor. + """ + return response + + async def post_list_models_with_metadata( + self, + response: model_service.ListModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_models_with_metadata` + interceptor in new development instead of the `post_list_models` interceptor. + When both interceptors are used, this `post_list_models_with_metadata` interceptor runs after the + `post_list_models` interceptor. The (possibly modified) response returned by + `post_list_models` will be passed to + `post_list_models_with_metadata`. + """ + return response, metadata + + async def pre_list_model_version_checkpoints( + self, + request: model_service.ListModelVersionCheckpointsRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to manipulate the request or metadata + before they are sent to the ModelService server. + """ + return request, metadata + + async def post_list_model_version_checkpoints( + self, response: model_service.ListModelVersionCheckpointsResponse + ) -> model_service.ListModelVersionCheckpointsResponse: + """Post-rpc interceptor for list_model_version_checkpoints + + DEPRECATED. Please use the `post_list_model_version_checkpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the ModelService server but before + it is returned to user code. This `post_list_model_version_checkpoints` interceptor runs + before the `post_list_model_version_checkpoints_with_metadata` interceptor. """ return response + async def post_list_model_version_checkpoints_with_metadata( + self, + response: model_service.ListModelVersionCheckpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_version_checkpoints_with_metadata` + interceptor in new development instead of the `post_list_model_version_checkpoints` interceptor. + When both interceptors are used, this `post_list_model_version_checkpoints_with_metadata` interceptor runs after the + `post_list_model_version_checkpoints` interceptor. The (possibly modified) response returned by + `post_list_model_version_checkpoints` will be passed to + `post_list_model_version_checkpoints_with_metadata`. + """ + return response, metadata + async def pre_list_model_versions( self, request: model_service.ListModelVersionsRequest, @@ -596,12 +974,37 @@ async def post_list_model_versions( ) -> model_service.ListModelVersionsResponse: """Post-rpc interceptor for list_model_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_versions` interceptor runs + before the `post_list_model_versions_with_metadata` interceptor. """ return response + async def post_list_model_versions_with_metadata( + self, + response: model_service.ListModelVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_model_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_versions_with_metadata` + interceptor in new development instead of the `post_list_model_versions` interceptor. + When both interceptors are used, this `post_list_model_versions_with_metadata` interceptor runs after the + `post_list_model_versions` interceptor. The (possibly modified) response returned by + `post_list_model_versions` will be passed to + `post_list_model_versions_with_metadata`. + """ + return response, metadata + async def pre_merge_version_aliases( self, request: model_service.MergeVersionAliasesRequest, @@ -620,12 +1023,33 @@ async def pre_merge_version_aliases( async def post_merge_version_aliases(self, response: model.Model) -> model.Model: """Post-rpc interceptor for merge_version_aliases - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_merge_version_aliases_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_merge_version_aliases` interceptor runs + before the `post_merge_version_aliases_with_metadata` interceptor. """ return response + async def post_merge_version_aliases_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for merge_version_aliases + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_merge_version_aliases_with_metadata` + interceptor in new development instead of the `post_merge_version_aliases` interceptor. + When both interceptors are used, this `post_merge_version_aliases_with_metadata` interceptor runs after the + `post_merge_version_aliases` interceptor. The (possibly modified) response returned by + `post_merge_version_aliases` will be passed to + `post_merge_version_aliases_with_metadata`. + """ + return response, metadata + async def pre_update_explanation_dataset( self, request: model_service.UpdateExplanationDatasetRequest, @@ -646,12 +1070,35 @@ async def post_update_explanation_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_explanation_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_explanation_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_explanation_dataset` interceptor runs + before the `post_update_explanation_dataset_with_metadata` interceptor. """ return response + async def post_update_explanation_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_explanation_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_explanation_dataset_with_metadata` + interceptor in new development instead of the `post_update_explanation_dataset` interceptor. + When both interceptors are used, this `post_update_explanation_dataset_with_metadata` interceptor runs after the + `post_update_explanation_dataset` interceptor. The (possibly modified) response returned by + `post_update_explanation_dataset` will be passed to + `post_update_explanation_dataset_with_metadata`. + """ + return response, metadata + async def pre_update_model( self, request: model_service.UpdateModelRequest, @@ -669,12 +1116,35 @@ async def pre_update_model( async def post_update_model(self, response: gca_model.Model) -> gca_model.Model: """Post-rpc interceptor for update_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_model` interceptor runs + before the `post_update_model_with_metadata` interceptor. """ return response + async def post_update_model_with_metadata( + self, + response: gca_model.Model, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_model_with_metadata` + interceptor in new development instead of the `post_update_model` interceptor. + When both interceptors are used, this `post_update_model_with_metadata` interceptor runs after the + `post_update_model` interceptor. The (possibly modified) response returned by + `post_update_model` will be passed to + `post_update_model_with_metadata`. + """ + return response, metadata + async def pre_upload_model( self, request: model_service.UploadModelRequest, @@ -694,12 +1164,35 @@ async def post_upload_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for upload_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_upload_model` interceptor runs + before the `post_upload_model_with_metadata` interceptor. """ return response + async def post_upload_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upload_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_upload_model_with_metadata` + interceptor in new development instead of the `post_upload_model` interceptor. + When both interceptors are used, this `post_upload_model_with_metadata` interceptor runs after the + `post_upload_model` interceptor. The (possibly modified) response returned by + `post_upload_model` will be passed to + `post_upload_model_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1035,6 +1528,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: self._wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: self._wrap_method( self.update_model, default_timeout=None, @@ -1303,6 +1801,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_import_evaluated_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_import_evaluated_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1474,6 +1979,13 @@ async def __call__( resp = await self._interceptor.post_batch_import_model_evaluation_slices( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_import_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1641,6 +2153,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_copy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_copy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1798,6 +2314,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1954,6 +2474,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_model_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2119,6 +2643,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2271,6 +2799,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2428,6 +2960,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2586,6 +3122,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_evaluation_slice(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_model_evaluation_slice_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2749,6 +3292,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_import_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2904,6 +3454,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_evaluations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_model_evaluations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3063,6 +3617,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3223,6 +3784,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3249,6 +3814,174 @@ async def __call__( return resp + class _ListModelVersionCheckpoints( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints, + AsyncModelServiceRestStub, + ): + def __hash__(self): + return hash("AsyncModelServiceRestTransport.ListModelVersionCheckpoints") + + @staticmethod + async def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = await getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + ) + return response + + async def __call__( + self, + request: model_service.ListModelVersionCheckpointsRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> model_service.ListModelVersionCheckpointsResponse: + r"""Call the list model version + checkpoints method over HTTP. + + Args: + request (~.model_service.ListModelVersionCheckpointsRequest): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints]. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.model_service.ListModelVersionCheckpointsResponse: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + + """ + + http_options = ( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_http_options() + ) + + ( + request, + metadata, + ) = await self._interceptor.pre_list_model_version_checkpoints( + request, metadata + ) + transcoded_request = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_transcoded_request( + http_options, request + ) + + # Jsonify the query params + query_params = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = type(request).to_json(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1.ModelServiceClient.ListModelVersionCheckpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = await AsyncModelServiceRestTransport._ListModelVersionCheckpoints._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + content = await response.read() + payload = json.loads(content.decode("utf-8")) + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + raise core_exceptions.format_http_response_error(response, method, request_url, payload) # type: ignore + + # Return the response + resp = model_service.ListModelVersionCheckpointsResponse() + pb_resp = model_service.ListModelVersionCheckpointsResponse.pb(resp) + content = await response.read() + json_format.Parse(content, pb_resp, ignore_unknown_fields=True) + resp = await self._interceptor.post_list_model_version_checkpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_version_checkpoints_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = ( + model_service.ListModelVersionCheckpointsResponse.to_json( + response + ) + ) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": "OK", # need to obtain this properly + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1.ModelServiceAsyncClient.list_model_version_checkpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + + return resp + class _ListModelVersions( _BaseModelServiceRestTransport._BaseListModelVersions, AsyncModelServiceRestStub ): @@ -3379,6 +4112,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_model_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3540,6 +4277,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_merge_version_aliases(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_merge_version_aliases_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3701,6 +4442,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_explanation_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_explanation_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3863,6 +4611,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4028,6 +4780,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upload_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upload_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6061,6 +6817,15 @@ def list_models( ) -> Callable[[model_service.ListModelsRequest], model_service.ListModelsResponse]: return self._ListModels(self._session, self._host, self._interceptor) # type: ignore + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + return self._ListModelVersionCheckpoints(self._session, self._host, self._interceptor) # type: ignore + @property def list_model_versions( self, diff --git a/google/cloud/aiplatform_v1/services/model_service/transports/rest_base.py b/google/cloud/aiplatform_v1/services/model_service/transports/rest_base.py index b8552150ae..dc4cdd4294 100644 --- a/google/cloud/aiplatform_v1/services/model_service/transports/rest_base.py +++ b/google/cloud/aiplatform_v1/services/model_service/transports/rest_base.py @@ -763,6 +763,53 @@ def _get_query_params_json(transcoded_request): query_params["$alt"] = "json;enum-encoding=int" return query_params + class _BaseListModelVersionCheckpoints: + def __hash__(self): # pragma: NO COVER + return NotImplementedError("__hash__ must be implemented.") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {} + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + @staticmethod + def _get_http_options(): + http_options: List[Dict[str, str]] = [ + { + "method": "get", + "uri": "/v1/{name=projects/*/locations/*/models/*}:listCheckpoints", + }, + ] + return http_options + + @staticmethod + def _get_transcoded_request(http_options, request): + pb_request = model_service.ListModelVersionCheckpointsRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + return transcoded_request + + @staticmethod + def _get_query_params_json(transcoded_request): + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_unset_required_fields( + query_params + ) + ) + + query_params["$alt"] = "json;enum-encoding=int" + return query_params + class _BaseListModelVersions: def __hash__(self): # pragma: NO COVER return NotImplementedError("__hash__ must be implemented.") diff --git a/google/cloud/aiplatform_v1/services/notebook_service/async_client.py b/google/cloud/aiplatform_v1/services/notebook_service/async_client.py index 646941b048..3a808ca7c8 100644 --- a/google/cloud/aiplatform_v1/services/notebook_service/async_client.py +++ b/google/cloud/aiplatform_v1/services/notebook_service/async_client.py @@ -61,6 +61,7 @@ from google.cloud.aiplatform_v1.types import notebook_runtime as gca_notebook_runtime from google.cloud.aiplatform_v1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1.types import notebook_service +from google.cloud.aiplatform_v1.types import notebook_software_config from google.cloud.aiplatform_v1.types import operation as gca_operation from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore diff --git a/google/cloud/aiplatform_v1/services/notebook_service/client.py b/google/cloud/aiplatform_v1/services/notebook_service/client.py index 32e26cd961..7aa92bfe8c 100644 --- a/google/cloud/aiplatform_v1/services/notebook_service/client.py +++ b/google/cloud/aiplatform_v1/services/notebook_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -75,6 +77,7 @@ from google.cloud.aiplatform_v1.types import notebook_runtime as gca_notebook_runtime from google.cloud.aiplatform_v1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1.types import notebook_service +from google.cloud.aiplatform_v1.types import notebook_software_config from google.cloud.aiplatform_v1.types import operation as gca_operation from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore @@ -653,6 +656,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3042,16 +3072,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3097,16 +3131,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3269,16 +3307,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3390,16 +3432,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3512,16 +3558,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3572,16 +3622,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3627,16 +3681,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3682,16 +3740,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/notebook_service/transports/rest.py b/google/cloud/aiplatform_v1/services/notebook_service/transports/rest.py index fb7fa60586..66c03d446f 100644 --- a/google/cloud/aiplatform_v1/services/notebook_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/notebook_service/transports/rest.py @@ -235,12 +235,35 @@ def post_assign_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for assign_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_assign_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_assign_notebook_runtime` interceptor runs + before the `post_assign_notebook_runtime_with_metadata` interceptor. """ return response + def post_assign_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for assign_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_assign_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_assign_notebook_runtime` interceptor. + When both interceptors are used, this `post_assign_notebook_runtime_with_metadata` interceptor runs after the + `post_assign_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_assign_notebook_runtime` will be passed to + `post_assign_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_create_notebook_execution_job( self, request: notebook_service.CreateNotebookExecutionJobRequest, @@ -261,12 +284,35 @@ def post_create_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_execution_job` interceptor runs + before the `post_create_notebook_execution_job_with_metadata` interceptor. """ return response + def post_create_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_create_notebook_execution_job` interceptor. + When both interceptors are used, this `post_create_notebook_execution_job_with_metadata` interceptor runs after the + `post_create_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_create_notebook_execution_job` will be passed to + `post_create_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_create_notebook_runtime_template( self, request: notebook_service.CreateNotebookRuntimeTemplateRequest, @@ -287,12 +333,35 @@ def post_create_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_runtime_template` interceptor runs + before the `post_create_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_create_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_create_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_create_notebook_runtime_template_with_metadata` interceptor runs after the + `post_create_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_create_notebook_runtime_template` will be passed to + `post_create_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_execution_job( self, request: notebook_service.DeleteNotebookExecutionJobRequest, @@ -313,12 +382,35 @@ def post_delete_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_execution_job` interceptor runs + before the `post_delete_notebook_execution_job_with_metadata` interceptor. """ return response + def post_delete_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_delete_notebook_execution_job` interceptor. + When both interceptors are used, this `post_delete_notebook_execution_job_with_metadata` interceptor runs after the + `post_delete_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_delete_notebook_execution_job` will be passed to + `post_delete_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_runtime( self, request: notebook_service.DeleteNotebookRuntimeRequest, @@ -339,12 +431,35 @@ def post_delete_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime` interceptor runs + before the `post_delete_notebook_runtime_with_metadata` interceptor. """ return response + def post_delete_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_with_metadata` interceptor runs after the + `post_delete_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime` will be passed to + `post_delete_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_runtime_template( self, request: notebook_service.DeleteNotebookRuntimeTemplateRequest, @@ -365,12 +480,35 @@ def post_delete_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime_template` interceptor runs + before the `post_delete_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_delete_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_template_with_metadata` interceptor runs after the + `post_delete_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime_template` will be passed to + `post_delete_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_get_notebook_execution_job( self, request: notebook_service.GetNotebookExecutionJobRequest, @@ -391,12 +529,38 @@ def post_get_notebook_execution_job( ) -> notebook_execution_job.NotebookExecutionJob: """Post-rpc interceptor for get_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_execution_job` interceptor runs + before the `post_get_notebook_execution_job_with_metadata` interceptor. """ return response + def post_get_notebook_execution_job_with_metadata( + self, + response: notebook_execution_job.NotebookExecutionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_execution_job.NotebookExecutionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_get_notebook_execution_job` interceptor. + When both interceptors are used, this `post_get_notebook_execution_job_with_metadata` interceptor runs after the + `post_get_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_get_notebook_execution_job` will be passed to + `post_get_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_get_notebook_runtime( self, request: notebook_service.GetNotebookRuntimeRequest, @@ -417,12 +581,37 @@ def post_get_notebook_runtime( ) -> notebook_runtime.NotebookRuntime: """Post-rpc interceptor for get_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime` interceptor runs + before the `post_get_notebook_runtime_with_metadata` interceptor. """ return response + def post_get_notebook_runtime_with_metadata( + self, + response: notebook_runtime.NotebookRuntime, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntime, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_with_metadata` interceptor runs after the + `post_get_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime` will be passed to + `post_get_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_get_notebook_runtime_template( self, request: notebook_service.GetNotebookRuntimeTemplateRequest, @@ -443,12 +632,38 @@ def post_get_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for get_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime_template` interceptor runs + before the `post_get_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_get_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_template_with_metadata` interceptor runs after the + `post_get_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime_template` will be passed to + `post_get_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_list_notebook_execution_jobs( self, request: notebook_service.ListNotebookExecutionJobsRequest, @@ -469,12 +684,38 @@ def post_list_notebook_execution_jobs( ) -> notebook_service.ListNotebookExecutionJobsResponse: """Post-rpc interceptor for list_notebook_execution_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_execution_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_execution_jobs` interceptor runs + before the `post_list_notebook_execution_jobs_with_metadata` interceptor. """ return response + def post_list_notebook_execution_jobs_with_metadata( + self, + response: notebook_service.ListNotebookExecutionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookExecutionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_execution_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_execution_jobs_with_metadata` + interceptor in new development instead of the `post_list_notebook_execution_jobs` interceptor. + When both interceptors are used, this `post_list_notebook_execution_jobs_with_metadata` interceptor runs after the + `post_list_notebook_execution_jobs` interceptor. The (possibly modified) response returned by + `post_list_notebook_execution_jobs` will be passed to + `post_list_notebook_execution_jobs_with_metadata`. + """ + return response, metadata + def pre_list_notebook_runtimes( self, request: notebook_service.ListNotebookRuntimesRequest, @@ -495,12 +736,38 @@ def post_list_notebook_runtimes( ) -> notebook_service.ListNotebookRuntimesResponse: """Post-rpc interceptor for list_notebook_runtimes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtimes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtimes` interceptor runs + before the `post_list_notebook_runtimes_with_metadata` interceptor. """ return response + def post_list_notebook_runtimes_with_metadata( + self, + response: notebook_service.ListNotebookRuntimesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtimes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtimes_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtimes` interceptor. + When both interceptors are used, this `post_list_notebook_runtimes_with_metadata` interceptor runs after the + `post_list_notebook_runtimes` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtimes` will be passed to + `post_list_notebook_runtimes_with_metadata`. + """ + return response, metadata + def pre_list_notebook_runtime_templates( self, request: notebook_service.ListNotebookRuntimeTemplatesRequest, @@ -521,12 +788,38 @@ def post_list_notebook_runtime_templates( ) -> notebook_service.ListNotebookRuntimeTemplatesResponse: """Post-rpc interceptor for list_notebook_runtime_templates - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtime_templates_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtime_templates` interceptor runs + before the `post_list_notebook_runtime_templates_with_metadata` interceptor. """ return response + def post_list_notebook_runtime_templates_with_metadata( + self, + response: notebook_service.ListNotebookRuntimeTemplatesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimeTemplatesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtime_templates + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtime_templates_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtime_templates` interceptor. + When both interceptors are used, this `post_list_notebook_runtime_templates_with_metadata` interceptor runs after the + `post_list_notebook_runtime_templates` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtime_templates` will be passed to + `post_list_notebook_runtime_templates_with_metadata`. + """ + return response, metadata + def pre_start_notebook_runtime( self, request: notebook_service.StartNotebookRuntimeRequest, @@ -547,12 +840,35 @@ def post_start_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for start_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_start_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_start_notebook_runtime` interceptor runs + before the `post_start_notebook_runtime_with_metadata` interceptor. """ return response + def post_start_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for start_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_start_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_start_notebook_runtime` interceptor. + When both interceptors are used, this `post_start_notebook_runtime_with_metadata` interceptor runs after the + `post_start_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_start_notebook_runtime` will be passed to + `post_start_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_stop_notebook_runtime( self, request: notebook_service.StopNotebookRuntimeRequest, @@ -573,12 +889,35 @@ def post_stop_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for stop_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_stop_notebook_runtime` interceptor runs + before the `post_stop_notebook_runtime_with_metadata` interceptor. """ return response + def post_stop_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_stop_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_stop_notebook_runtime` interceptor. + When both interceptors are used, this `post_stop_notebook_runtime_with_metadata` interceptor runs after the + `post_stop_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_stop_notebook_runtime` will be passed to + `post_stop_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_update_notebook_runtime_template( self, request: notebook_service.UpdateNotebookRuntimeTemplateRequest, @@ -599,12 +938,38 @@ def post_update_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for update_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_update_notebook_runtime_template` interceptor runs + before the `post_update_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_update_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_update_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_update_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_update_notebook_runtime_template_with_metadata` interceptor runs after the + `post_update_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_update_notebook_runtime_template` will be passed to + `post_update_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_upgrade_notebook_runtime( self, request: notebook_service.UpgradeNotebookRuntimeRequest, @@ -625,12 +990,35 @@ def post_upgrade_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for upgrade_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upgrade_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_upgrade_notebook_runtime` interceptor runs + before the `post_upgrade_notebook_runtime_with_metadata` interceptor. """ return response + def post_upgrade_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upgrade_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_upgrade_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_upgrade_notebook_runtime` interceptor. + When both interceptors are used, this `post_upgrade_notebook_runtime_with_metadata` interceptor runs after the + `post_upgrade_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_upgrade_notebook_runtime` will be passed to + `post_upgrade_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3008,6 +3396,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_assign_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_assign_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3163,6 +3555,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3316,6 +3715,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3465,6 +3871,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3613,6 +4026,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3760,6 +4177,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3910,6 +4334,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4061,6 +4489,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4216,6 +4648,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4368,6 +4807,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_execution_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_notebook_execution_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4519,6 +4962,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_runtimes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_notebook_runtimes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4669,6 +5116,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_runtime_templates(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_notebook_runtime_templates_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4825,6 +5279,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_start_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_start_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4977,6 +5435,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_stop_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stop_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5134,6 +5596,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5290,6 +5759,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_upgrade_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upgrade_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/notebook_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/notebook_service/transports/rest_asyncio.py index 0a2cf729d7..827c472097 100644 --- a/google/cloud/aiplatform_v1/services/notebook_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/notebook_service/transports/rest_asyncio.py @@ -253,12 +253,35 @@ async def post_assign_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for assign_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_assign_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_assign_notebook_runtime` interceptor runs + before the `post_assign_notebook_runtime_with_metadata` interceptor. """ return response + async def post_assign_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for assign_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_assign_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_assign_notebook_runtime` interceptor. + When both interceptors are used, this `post_assign_notebook_runtime_with_metadata` interceptor runs after the + `post_assign_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_assign_notebook_runtime` will be passed to + `post_assign_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_create_notebook_execution_job( self, request: notebook_service.CreateNotebookExecutionJobRequest, @@ -279,12 +302,35 @@ async def post_create_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_execution_job` interceptor runs + before the `post_create_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_create_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_create_notebook_execution_job` interceptor. + When both interceptors are used, this `post_create_notebook_execution_job_with_metadata` interceptor runs after the + `post_create_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_create_notebook_execution_job` will be passed to + `post_create_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_create_notebook_runtime_template( self, request: notebook_service.CreateNotebookRuntimeTemplateRequest, @@ -305,12 +351,35 @@ async def post_create_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_runtime_template` interceptor runs + before the `post_create_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_create_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_create_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_create_notebook_runtime_template_with_metadata` interceptor runs after the + `post_create_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_create_notebook_runtime_template` will be passed to + `post_create_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_execution_job( self, request: notebook_service.DeleteNotebookExecutionJobRequest, @@ -331,12 +400,35 @@ async def post_delete_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_execution_job` interceptor runs + before the `post_delete_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_delete_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_delete_notebook_execution_job` interceptor. + When both interceptors are used, this `post_delete_notebook_execution_job_with_metadata` interceptor runs after the + `post_delete_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_delete_notebook_execution_job` will be passed to + `post_delete_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_runtime( self, request: notebook_service.DeleteNotebookRuntimeRequest, @@ -357,12 +449,35 @@ async def post_delete_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime` interceptor runs + before the `post_delete_notebook_runtime_with_metadata` interceptor. """ return response + async def post_delete_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_with_metadata` interceptor runs after the + `post_delete_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime` will be passed to + `post_delete_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_runtime_template( self, request: notebook_service.DeleteNotebookRuntimeTemplateRequest, @@ -383,12 +498,35 @@ async def post_delete_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime_template` interceptor runs + before the `post_delete_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_delete_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_template_with_metadata` interceptor runs after the + `post_delete_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime_template` will be passed to + `post_delete_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_execution_job( self, request: notebook_service.GetNotebookExecutionJobRequest, @@ -409,12 +547,38 @@ async def post_get_notebook_execution_job( ) -> notebook_execution_job.NotebookExecutionJob: """Post-rpc interceptor for get_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_execution_job` interceptor runs + before the `post_get_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_get_notebook_execution_job_with_metadata( + self, + response: notebook_execution_job.NotebookExecutionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_execution_job.NotebookExecutionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_get_notebook_execution_job` interceptor. + When both interceptors are used, this `post_get_notebook_execution_job_with_metadata` interceptor runs after the + `post_get_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_get_notebook_execution_job` will be passed to + `post_get_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_runtime( self, request: notebook_service.GetNotebookRuntimeRequest, @@ -435,12 +599,37 @@ async def post_get_notebook_runtime( ) -> notebook_runtime.NotebookRuntime: """Post-rpc interceptor for get_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime` interceptor runs + before the `post_get_notebook_runtime_with_metadata` interceptor. """ return response + async def post_get_notebook_runtime_with_metadata( + self, + response: notebook_runtime.NotebookRuntime, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntime, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_with_metadata` interceptor runs after the + `post_get_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime` will be passed to + `post_get_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_runtime_template( self, request: notebook_service.GetNotebookRuntimeTemplateRequest, @@ -461,12 +650,38 @@ async def post_get_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for get_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime_template` interceptor runs + before the `post_get_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_get_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_template_with_metadata` interceptor runs after the + `post_get_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime_template` will be passed to + `post_get_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_execution_jobs( self, request: notebook_service.ListNotebookExecutionJobsRequest, @@ -487,12 +702,38 @@ async def post_list_notebook_execution_jobs( ) -> notebook_service.ListNotebookExecutionJobsResponse: """Post-rpc interceptor for list_notebook_execution_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_execution_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_execution_jobs` interceptor runs + before the `post_list_notebook_execution_jobs_with_metadata` interceptor. """ return response + async def post_list_notebook_execution_jobs_with_metadata( + self, + response: notebook_service.ListNotebookExecutionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookExecutionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_execution_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_execution_jobs_with_metadata` + interceptor in new development instead of the `post_list_notebook_execution_jobs` interceptor. + When both interceptors are used, this `post_list_notebook_execution_jobs_with_metadata` interceptor runs after the + `post_list_notebook_execution_jobs` interceptor. The (possibly modified) response returned by + `post_list_notebook_execution_jobs` will be passed to + `post_list_notebook_execution_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_runtimes( self, request: notebook_service.ListNotebookRuntimesRequest, @@ -513,12 +754,38 @@ async def post_list_notebook_runtimes( ) -> notebook_service.ListNotebookRuntimesResponse: """Post-rpc interceptor for list_notebook_runtimes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtimes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtimes` interceptor runs + before the `post_list_notebook_runtimes_with_metadata` interceptor. """ return response + async def post_list_notebook_runtimes_with_metadata( + self, + response: notebook_service.ListNotebookRuntimesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtimes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtimes_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtimes` interceptor. + When both interceptors are used, this `post_list_notebook_runtimes_with_metadata` interceptor runs after the + `post_list_notebook_runtimes` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtimes` will be passed to + `post_list_notebook_runtimes_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_runtime_templates( self, request: notebook_service.ListNotebookRuntimeTemplatesRequest, @@ -539,12 +806,38 @@ async def post_list_notebook_runtime_templates( ) -> notebook_service.ListNotebookRuntimeTemplatesResponse: """Post-rpc interceptor for list_notebook_runtime_templates - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtime_templates_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtime_templates` interceptor runs + before the `post_list_notebook_runtime_templates_with_metadata` interceptor. """ return response + async def post_list_notebook_runtime_templates_with_metadata( + self, + response: notebook_service.ListNotebookRuntimeTemplatesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimeTemplatesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtime_templates + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtime_templates_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtime_templates` interceptor. + When both interceptors are used, this `post_list_notebook_runtime_templates_with_metadata` interceptor runs after the + `post_list_notebook_runtime_templates` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtime_templates` will be passed to + `post_list_notebook_runtime_templates_with_metadata`. + """ + return response, metadata + async def pre_start_notebook_runtime( self, request: notebook_service.StartNotebookRuntimeRequest, @@ -565,12 +858,35 @@ async def post_start_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for start_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_start_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_start_notebook_runtime` interceptor runs + before the `post_start_notebook_runtime_with_metadata` interceptor. """ return response + async def post_start_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for start_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_start_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_start_notebook_runtime` interceptor. + When both interceptors are used, this `post_start_notebook_runtime_with_metadata` interceptor runs after the + `post_start_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_start_notebook_runtime` will be passed to + `post_start_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_stop_notebook_runtime( self, request: notebook_service.StopNotebookRuntimeRequest, @@ -591,12 +907,35 @@ async def post_stop_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for stop_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_stop_notebook_runtime` interceptor runs + before the `post_stop_notebook_runtime_with_metadata` interceptor. """ return response + async def post_stop_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_stop_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_stop_notebook_runtime` interceptor. + When both interceptors are used, this `post_stop_notebook_runtime_with_metadata` interceptor runs after the + `post_stop_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_stop_notebook_runtime` will be passed to + `post_stop_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_update_notebook_runtime_template( self, request: notebook_service.UpdateNotebookRuntimeTemplateRequest, @@ -617,12 +956,38 @@ async def post_update_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for update_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_update_notebook_runtime_template` interceptor runs + before the `post_update_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_update_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_update_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_update_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_update_notebook_runtime_template_with_metadata` interceptor runs after the + `post_update_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_update_notebook_runtime_template` will be passed to + `post_update_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_upgrade_notebook_runtime( self, request: notebook_service.UpgradeNotebookRuntimeRequest, @@ -643,12 +1008,35 @@ async def post_upgrade_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for upgrade_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upgrade_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_upgrade_notebook_runtime` interceptor runs + before the `post_upgrade_notebook_runtime_with_metadata` interceptor. """ return response + async def post_upgrade_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upgrade_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_upgrade_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_upgrade_notebook_runtime` interceptor. + When both interceptors are used, this `post_upgrade_notebook_runtime_with_metadata` interceptor runs after the + `post_upgrade_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_upgrade_notebook_runtime` will be passed to + `post_upgrade_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1238,6 +1626,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_assign_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_assign_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1402,6 +1797,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1568,6 +1970,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1726,6 +2135,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1880,6 +2296,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2040,6 +2463,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2194,6 +2624,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2351,6 +2788,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2513,6 +2954,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2672,6 +3120,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_execution_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_notebook_execution_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2829,6 +3284,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_runtimes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_notebook_runtimes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2990,6 +3449,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_runtime_templates(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_notebook_runtime_templates_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3154,6 +3620,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_start_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_start_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3314,6 +3784,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_stop_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stop_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3482,6 +3956,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3644,6 +4125,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upgrade_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_upgrade_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/persistent_resource_service/client.py b/google/cloud/aiplatform_v1/services/persistent_resource_service/client.py index ab694c7173..fbb7673ec3 100644 --- a/google/cloud/aiplatform_v1/services/persistent_resource_service/client.py +++ b/google/cloud/aiplatform_v1/services/persistent_resource_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -560,6 +562,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1658,16 +1687,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1713,16 +1746,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1885,16 +1922,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2006,16 +2047,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2128,16 +2173,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2188,16 +2237,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2243,16 +2296,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2298,16 +2355,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest.py b/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest.py index d9a47ed2a5..06af53f1c6 100644 --- a/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest.py @@ -154,12 +154,35 @@ def post_create_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_create_persistent_resource` interceptor runs + before the `post_create_persistent_resource_with_metadata` interceptor. """ return response + def post_create_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_create_persistent_resource_with_metadata` + interceptor in new development instead of the `post_create_persistent_resource` interceptor. + When both interceptors are used, this `post_create_persistent_resource_with_metadata` interceptor runs after the + `post_create_persistent_resource` interceptor. The (possibly modified) response returned by + `post_create_persistent_resource` will be passed to + `post_create_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_delete_persistent_resource( self, request: persistent_resource_service.DeletePersistentResourceRequest, @@ -180,12 +203,35 @@ def post_delete_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_delete_persistent_resource` interceptor runs + before the `post_delete_persistent_resource_with_metadata` interceptor. """ return response + def post_delete_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_delete_persistent_resource_with_metadata` + interceptor in new development instead of the `post_delete_persistent_resource` interceptor. + When both interceptors are used, this `post_delete_persistent_resource_with_metadata` interceptor runs after the + `post_delete_persistent_resource` interceptor. The (possibly modified) response returned by + `post_delete_persistent_resource` will be passed to + `post_delete_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_get_persistent_resource( self, request: persistent_resource_service.GetPersistentResourceRequest, @@ -206,12 +252,37 @@ def post_get_persistent_resource( ) -> persistent_resource.PersistentResource: """Post-rpc interceptor for get_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_get_persistent_resource` interceptor runs + before the `post_get_persistent_resource_with_metadata` interceptor. """ return response + def post_get_persistent_resource_with_metadata( + self, + response: persistent_resource.PersistentResource, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource.PersistentResource, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_get_persistent_resource_with_metadata` + interceptor in new development instead of the `post_get_persistent_resource` interceptor. + When both interceptors are used, this `post_get_persistent_resource_with_metadata` interceptor runs after the + `post_get_persistent_resource` interceptor. The (possibly modified) response returned by + `post_get_persistent_resource` will be passed to + `post_get_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_list_persistent_resources( self, request: persistent_resource_service.ListPersistentResourcesRequest, @@ -232,12 +303,38 @@ def post_list_persistent_resources( ) -> persistent_resource_service.ListPersistentResourcesResponse: """Post-rpc interceptor for list_persistent_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_persistent_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_list_persistent_resources` interceptor runs + before the `post_list_persistent_resources_with_metadata` interceptor. """ return response + def post_list_persistent_resources_with_metadata( + self, + response: persistent_resource_service.ListPersistentResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource_service.ListPersistentResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_persistent_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_list_persistent_resources_with_metadata` + interceptor in new development instead of the `post_list_persistent_resources` interceptor. + When both interceptors are used, this `post_list_persistent_resources_with_metadata` interceptor runs after the + `post_list_persistent_resources` interceptor. The (possibly modified) response returned by + `post_list_persistent_resources` will be passed to + `post_list_persistent_resources_with_metadata`. + """ + return response, metadata + def pre_reboot_persistent_resource( self, request: persistent_resource_service.RebootPersistentResourceRequest, @@ -258,12 +355,35 @@ def post_reboot_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for reboot_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_reboot_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_reboot_persistent_resource` interceptor runs + before the `post_reboot_persistent_resource_with_metadata` interceptor. """ return response + def post_reboot_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for reboot_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_reboot_persistent_resource_with_metadata` + interceptor in new development instead of the `post_reboot_persistent_resource` interceptor. + When both interceptors are used, this `post_reboot_persistent_resource_with_metadata` interceptor runs after the + `post_reboot_persistent_resource` interceptor. The (possibly modified) response returned by + `post_reboot_persistent_resource` will be passed to + `post_reboot_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_update_persistent_resource( self, request: persistent_resource_service.UpdatePersistentResourceRequest, @@ -284,12 +404,35 @@ def post_update_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_update_persistent_resource` interceptor runs + before the `post_update_persistent_resource_with_metadata` interceptor. """ return response + def post_update_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_update_persistent_resource_with_metadata` + interceptor in new development instead of the `post_update_persistent_resource` interceptor. + When both interceptors are used, this `post_update_persistent_resource_with_metadata` interceptor runs after the + `post_update_persistent_resource` interceptor. The (possibly modified) response returned by + `post_update_persistent_resource` will be passed to + `post_update_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2670,6 +2813,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2819,6 +2966,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2969,6 +3120,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3122,6 +3277,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_persistent_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_persistent_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3279,6 +3438,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_reboot_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_reboot_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3434,6 +3597,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest_asyncio.py index 5553e4a51d..fc0175da55 100644 --- a/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/persistent_resource_service/transports/rest_asyncio.py @@ -172,12 +172,35 @@ async def post_create_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_create_persistent_resource` interceptor runs + before the `post_create_persistent_resource_with_metadata` interceptor. """ return response + async def post_create_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_create_persistent_resource_with_metadata` + interceptor in new development instead of the `post_create_persistent_resource` interceptor. + When both interceptors are used, this `post_create_persistent_resource_with_metadata` interceptor runs after the + `post_create_persistent_resource` interceptor. The (possibly modified) response returned by + `post_create_persistent_resource` will be passed to + `post_create_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_delete_persistent_resource( self, request: persistent_resource_service.DeletePersistentResourceRequest, @@ -198,12 +221,35 @@ async def post_delete_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_delete_persistent_resource` interceptor runs + before the `post_delete_persistent_resource_with_metadata` interceptor. """ return response + async def post_delete_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_delete_persistent_resource_with_metadata` + interceptor in new development instead of the `post_delete_persistent_resource` interceptor. + When both interceptors are used, this `post_delete_persistent_resource_with_metadata` interceptor runs after the + `post_delete_persistent_resource` interceptor. The (possibly modified) response returned by + `post_delete_persistent_resource` will be passed to + `post_delete_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_get_persistent_resource( self, request: persistent_resource_service.GetPersistentResourceRequest, @@ -224,12 +270,37 @@ async def post_get_persistent_resource( ) -> persistent_resource.PersistentResource: """Post-rpc interceptor for get_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_get_persistent_resource` interceptor runs + before the `post_get_persistent_resource_with_metadata` interceptor. """ return response + async def post_get_persistent_resource_with_metadata( + self, + response: persistent_resource.PersistentResource, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource.PersistentResource, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_get_persistent_resource_with_metadata` + interceptor in new development instead of the `post_get_persistent_resource` interceptor. + When both interceptors are used, this `post_get_persistent_resource_with_metadata` interceptor runs after the + `post_get_persistent_resource` interceptor. The (possibly modified) response returned by + `post_get_persistent_resource` will be passed to + `post_get_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_list_persistent_resources( self, request: persistent_resource_service.ListPersistentResourcesRequest, @@ -250,12 +321,38 @@ async def post_list_persistent_resources( ) -> persistent_resource_service.ListPersistentResourcesResponse: """Post-rpc interceptor for list_persistent_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_persistent_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_list_persistent_resources` interceptor runs + before the `post_list_persistent_resources_with_metadata` interceptor. """ return response + async def post_list_persistent_resources_with_metadata( + self, + response: persistent_resource_service.ListPersistentResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource_service.ListPersistentResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_persistent_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_list_persistent_resources_with_metadata` + interceptor in new development instead of the `post_list_persistent_resources` interceptor. + When both interceptors are used, this `post_list_persistent_resources_with_metadata` interceptor runs after the + `post_list_persistent_resources` interceptor. The (possibly modified) response returned by + `post_list_persistent_resources` will be passed to + `post_list_persistent_resources_with_metadata`. + """ + return response, metadata + async def pre_reboot_persistent_resource( self, request: persistent_resource_service.RebootPersistentResourceRequest, @@ -276,12 +373,35 @@ async def post_reboot_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for reboot_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_reboot_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_reboot_persistent_resource` interceptor runs + before the `post_reboot_persistent_resource_with_metadata` interceptor. """ return response + async def post_reboot_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for reboot_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_reboot_persistent_resource_with_metadata` + interceptor in new development instead of the `post_reboot_persistent_resource` interceptor. + When both interceptors are used, this `post_reboot_persistent_resource_with_metadata` interceptor runs after the + `post_reboot_persistent_resource` interceptor. The (possibly modified) response returned by + `post_reboot_persistent_resource` will be passed to + `post_reboot_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_update_persistent_resource( self, request: persistent_resource_service.UpdatePersistentResourceRequest, @@ -302,12 +422,35 @@ async def post_update_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_update_persistent_resource` interceptor runs + before the `post_update_persistent_resource_with_metadata` interceptor. """ return response + async def post_update_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_update_persistent_resource_with_metadata` + interceptor in new development instead of the `post_update_persistent_resource` interceptor. + When both interceptors are used, this `post_update_persistent_resource_with_metadata` interceptor runs after the + `post_update_persistent_resource` interceptor. The (possibly modified) response returned by + `post_update_persistent_resource` will be passed to + `post_update_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -854,6 +997,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1011,6 +1161,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1169,6 +1326,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1328,6 +1492,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_persistent_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_persistent_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1493,6 +1664,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_reboot_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_reboot_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1656,6 +1834,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/pipeline_service/client.py b/google/cloud/aiplatform_v1/services/pipeline_service/client.py index 7720e7f25a..e27c896ce3 100644 --- a/google/cloud/aiplatform_v1/services/pipeline_service/client.py +++ b/google/cloud/aiplatform_v1/services/pipeline_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -698,6 +700,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2457,16 +2486,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2512,16 +2545,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2684,16 +2721,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2805,16 +2846,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2927,16 +2972,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2987,16 +3036,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3042,16 +3095,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3097,16 +3154,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest.py b/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest.py index d10267d680..edb0abb0bf 100644 --- a/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest.py @@ -198,12 +198,35 @@ def post_batch_cancel_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_cancel_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_cancel_pipeline_jobs` interceptor runs + before the `post_batch_cancel_pipeline_jobs_with_metadata` interceptor. """ return response + def post_batch_cancel_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_cancel_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_cancel_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_cancel_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_cancel_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_cancel_pipeline_jobs` will be passed to + `post_batch_cancel_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_batch_delete_pipeline_jobs( self, request: pipeline_service.BatchDeletePipelineJobsRequest, @@ -224,12 +247,35 @@ def post_batch_delete_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_delete_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_delete_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_delete_pipeline_jobs` interceptor runs + before the `post_batch_delete_pipeline_jobs_with_metadata` interceptor. """ return response + def post_batch_delete_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_delete_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_delete_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_delete_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_delete_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_delete_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_delete_pipeline_jobs` will be passed to + `post_batch_delete_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_cancel_pipeline_job( self, request: pipeline_service.CancelPipelineJobRequest, @@ -280,12 +326,35 @@ def post_create_pipeline_job( ) -> gca_pipeline_job.PipelineJob: """Post-rpc interceptor for create_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_pipeline_job` interceptor runs + before the `post_create_pipeline_job_with_metadata` interceptor. """ return response + def post_create_pipeline_job_with_metadata( + self, + response: gca_pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_pipeline_job_with_metadata` + interceptor in new development instead of the `post_create_pipeline_job` interceptor. + When both interceptors are used, this `post_create_pipeline_job_with_metadata` interceptor runs after the + `post_create_pipeline_job` interceptor. The (possibly modified) response returned by + `post_create_pipeline_job` will be passed to + `post_create_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_create_training_pipeline( self, request: pipeline_service.CreateTrainingPipelineRequest, @@ -306,12 +375,37 @@ def post_create_training_pipeline( ) -> gca_training_pipeline.TrainingPipeline: """Post-rpc interceptor for create_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_training_pipeline` interceptor runs + before the `post_create_training_pipeline_with_metadata` interceptor. """ return response + def post_create_training_pipeline_with_metadata( + self, + response: gca_training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_training_pipeline_with_metadata` + interceptor in new development instead of the `post_create_training_pipeline` interceptor. + When both interceptors are used, this `post_create_training_pipeline_with_metadata` interceptor runs after the + `post_create_training_pipeline` interceptor. The (possibly modified) response returned by + `post_create_training_pipeline` will be passed to + `post_create_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_delete_pipeline_job( self, request: pipeline_service.DeletePipelineJobRequest, @@ -332,12 +426,35 @@ def post_delete_pipeline_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_pipeline_job` interceptor runs + before the `post_delete_pipeline_job_with_metadata` interceptor. """ return response + def post_delete_pipeline_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_pipeline_job_with_metadata` + interceptor in new development instead of the `post_delete_pipeline_job` interceptor. + When both interceptors are used, this `post_delete_pipeline_job_with_metadata` interceptor runs after the + `post_delete_pipeline_job` interceptor. The (possibly modified) response returned by + `post_delete_pipeline_job` will be passed to + `post_delete_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_delete_training_pipeline( self, request: pipeline_service.DeleteTrainingPipelineRequest, @@ -358,12 +475,35 @@ def post_delete_training_pipeline( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_training_pipeline` interceptor runs + before the `post_delete_training_pipeline_with_metadata` interceptor. """ return response + def post_delete_training_pipeline_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_training_pipeline_with_metadata` + interceptor in new development instead of the `post_delete_training_pipeline` interceptor. + When both interceptors are used, this `post_delete_training_pipeline_with_metadata` interceptor runs after the + `post_delete_training_pipeline` interceptor. The (possibly modified) response returned by + `post_delete_training_pipeline` will be passed to + `post_delete_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_get_pipeline_job( self, request: pipeline_service.GetPipelineJobRequest, @@ -383,12 +523,35 @@ def post_get_pipeline_job( ) -> pipeline_job.PipelineJob: """Post-rpc interceptor for get_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_pipeline_job` interceptor runs + before the `post_get_pipeline_job_with_metadata` interceptor. """ return response + def post_get_pipeline_job_with_metadata( + self, + response: pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_pipeline_job_with_metadata` + interceptor in new development instead of the `post_get_pipeline_job` interceptor. + When both interceptors are used, this `post_get_pipeline_job_with_metadata` interceptor runs after the + `post_get_pipeline_job` interceptor. The (possibly modified) response returned by + `post_get_pipeline_job` will be passed to + `post_get_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_get_training_pipeline( self, request: pipeline_service.GetTrainingPipelineRequest, @@ -409,12 +572,37 @@ def post_get_training_pipeline( ) -> training_pipeline.TrainingPipeline: """Post-rpc interceptor for get_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_training_pipeline` interceptor runs + before the `post_get_training_pipeline_with_metadata` interceptor. """ return response + def post_get_training_pipeline_with_metadata( + self, + response: training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_training_pipeline_with_metadata` + interceptor in new development instead of the `post_get_training_pipeline` interceptor. + When both interceptors are used, this `post_get_training_pipeline_with_metadata` interceptor runs after the + `post_get_training_pipeline` interceptor. The (possibly modified) response returned by + `post_get_training_pipeline` will be passed to + `post_get_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_list_pipeline_jobs( self, request: pipeline_service.ListPipelineJobsRequest, @@ -435,12 +623,38 @@ def post_list_pipeline_jobs( ) -> pipeline_service.ListPipelineJobsResponse: """Post-rpc interceptor for list_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_pipeline_jobs` interceptor runs + before the `post_list_pipeline_jobs_with_metadata` interceptor. """ return response + def post_list_pipeline_jobs_with_metadata( + self, + response: pipeline_service.ListPipelineJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListPipelineJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_list_pipeline_jobs` interceptor. + When both interceptors are used, this `post_list_pipeline_jobs_with_metadata` interceptor runs after the + `post_list_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_list_pipeline_jobs` will be passed to + `post_list_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_list_training_pipelines( self, request: pipeline_service.ListTrainingPipelinesRequest, @@ -461,12 +675,38 @@ def post_list_training_pipelines( ) -> pipeline_service.ListTrainingPipelinesResponse: """Post-rpc interceptor for list_training_pipelines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_training_pipelines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_training_pipelines` interceptor runs + before the `post_list_training_pipelines_with_metadata` interceptor. """ return response + def post_list_training_pipelines_with_metadata( + self, + response: pipeline_service.ListTrainingPipelinesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListTrainingPipelinesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_training_pipelines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_training_pipelines_with_metadata` + interceptor in new development instead of the `post_list_training_pipelines` interceptor. + When both interceptors are used, this `post_list_training_pipelines_with_metadata` interceptor runs after the + `post_list_training_pipelines` interceptor. The (possibly modified) response returned by + `post_list_training_pipelines` will be passed to + `post_list_training_pipelines_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2847,6 +3087,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_cancel_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_cancel_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3002,6 +3246,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_delete_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_delete_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3391,6 +3639,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3550,6 +3802,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3698,6 +3954,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3846,6 +4106,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3992,6 +4256,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4143,6 +4411,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4291,6 +4563,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4442,6 +4718,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_training_pipelines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_training_pipelines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest_asyncio.py index a9894be3af..f54c6fc4f3 100644 --- a/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/pipeline_service/transports/rest_asyncio.py @@ -216,12 +216,35 @@ async def post_batch_cancel_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_cancel_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_cancel_pipeline_jobs` interceptor runs + before the `post_batch_cancel_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_batch_cancel_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_cancel_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_cancel_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_cancel_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_cancel_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_cancel_pipeline_jobs` will be passed to + `post_batch_cancel_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_batch_delete_pipeline_jobs( self, request: pipeline_service.BatchDeletePipelineJobsRequest, @@ -242,12 +265,35 @@ async def post_batch_delete_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_delete_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_delete_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_delete_pipeline_jobs` interceptor runs + before the `post_batch_delete_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_batch_delete_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_delete_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_delete_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_delete_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_delete_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_delete_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_delete_pipeline_jobs` will be passed to + `post_batch_delete_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_cancel_pipeline_job( self, request: pipeline_service.CancelPipelineJobRequest, @@ -298,12 +344,35 @@ async def post_create_pipeline_job( ) -> gca_pipeline_job.PipelineJob: """Post-rpc interceptor for create_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_pipeline_job` interceptor runs + before the `post_create_pipeline_job_with_metadata` interceptor. """ return response + async def post_create_pipeline_job_with_metadata( + self, + response: gca_pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_pipeline_job_with_metadata` + interceptor in new development instead of the `post_create_pipeline_job` interceptor. + When both interceptors are used, this `post_create_pipeline_job_with_metadata` interceptor runs after the + `post_create_pipeline_job` interceptor. The (possibly modified) response returned by + `post_create_pipeline_job` will be passed to + `post_create_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_create_training_pipeline( self, request: pipeline_service.CreateTrainingPipelineRequest, @@ -324,12 +393,37 @@ async def post_create_training_pipeline( ) -> gca_training_pipeline.TrainingPipeline: """Post-rpc interceptor for create_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_training_pipeline` interceptor runs + before the `post_create_training_pipeline_with_metadata` interceptor. """ return response + async def post_create_training_pipeline_with_metadata( + self, + response: gca_training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_training_pipeline_with_metadata` + interceptor in new development instead of the `post_create_training_pipeline` interceptor. + When both interceptors are used, this `post_create_training_pipeline_with_metadata` interceptor runs after the + `post_create_training_pipeline` interceptor. The (possibly modified) response returned by + `post_create_training_pipeline` will be passed to + `post_create_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_delete_pipeline_job( self, request: pipeline_service.DeletePipelineJobRequest, @@ -350,12 +444,35 @@ async def post_delete_pipeline_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_pipeline_job` interceptor runs + before the `post_delete_pipeline_job_with_metadata` interceptor. """ return response + async def post_delete_pipeline_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_pipeline_job_with_metadata` + interceptor in new development instead of the `post_delete_pipeline_job` interceptor. + When both interceptors are used, this `post_delete_pipeline_job_with_metadata` interceptor runs after the + `post_delete_pipeline_job` interceptor. The (possibly modified) response returned by + `post_delete_pipeline_job` will be passed to + `post_delete_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_delete_training_pipeline( self, request: pipeline_service.DeleteTrainingPipelineRequest, @@ -376,12 +493,35 @@ async def post_delete_training_pipeline( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_training_pipeline` interceptor runs + before the `post_delete_training_pipeline_with_metadata` interceptor. """ return response + async def post_delete_training_pipeline_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_training_pipeline_with_metadata` + interceptor in new development instead of the `post_delete_training_pipeline` interceptor. + When both interceptors are used, this `post_delete_training_pipeline_with_metadata` interceptor runs after the + `post_delete_training_pipeline` interceptor. The (possibly modified) response returned by + `post_delete_training_pipeline` will be passed to + `post_delete_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_get_pipeline_job( self, request: pipeline_service.GetPipelineJobRequest, @@ -401,12 +541,35 @@ async def post_get_pipeline_job( ) -> pipeline_job.PipelineJob: """Post-rpc interceptor for get_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_pipeline_job` interceptor runs + before the `post_get_pipeline_job_with_metadata` interceptor. """ return response + async def post_get_pipeline_job_with_metadata( + self, + response: pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_pipeline_job_with_metadata` + interceptor in new development instead of the `post_get_pipeline_job` interceptor. + When both interceptors are used, this `post_get_pipeline_job_with_metadata` interceptor runs after the + `post_get_pipeline_job` interceptor. The (possibly modified) response returned by + `post_get_pipeline_job` will be passed to + `post_get_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_get_training_pipeline( self, request: pipeline_service.GetTrainingPipelineRequest, @@ -427,12 +590,37 @@ async def post_get_training_pipeline( ) -> training_pipeline.TrainingPipeline: """Post-rpc interceptor for get_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_training_pipeline` interceptor runs + before the `post_get_training_pipeline_with_metadata` interceptor. """ return response + async def post_get_training_pipeline_with_metadata( + self, + response: training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_training_pipeline_with_metadata` + interceptor in new development instead of the `post_get_training_pipeline` interceptor. + When both interceptors are used, this `post_get_training_pipeline_with_metadata` interceptor runs after the + `post_get_training_pipeline` interceptor. The (possibly modified) response returned by + `post_get_training_pipeline` will be passed to + `post_get_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_list_pipeline_jobs( self, request: pipeline_service.ListPipelineJobsRequest, @@ -453,12 +641,38 @@ async def post_list_pipeline_jobs( ) -> pipeline_service.ListPipelineJobsResponse: """Post-rpc interceptor for list_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_pipeline_jobs` interceptor runs + before the `post_list_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_list_pipeline_jobs_with_metadata( + self, + response: pipeline_service.ListPipelineJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListPipelineJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_list_pipeline_jobs` interceptor. + When both interceptors are used, this `post_list_pipeline_jobs_with_metadata` interceptor runs after the + `post_list_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_list_pipeline_jobs` will be passed to + `post_list_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_training_pipelines( self, request: pipeline_service.ListTrainingPipelinesRequest, @@ -479,12 +693,38 @@ async def post_list_training_pipelines( ) -> pipeline_service.ListTrainingPipelinesResponse: """Post-rpc interceptor for list_training_pipelines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_training_pipelines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_training_pipelines` interceptor runs + before the `post_list_training_pipelines_with_metadata` interceptor. """ return response + async def post_list_training_pipelines_with_metadata( + self, + response: pipeline_service.ListTrainingPipelinesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListTrainingPipelinesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_training_pipelines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_training_pipelines_with_metadata` + interceptor in new development instead of the `post_list_training_pipelines` interceptor. + When both interceptors are used, this `post_list_training_pipelines_with_metadata` interceptor runs after the + `post_list_training_pipelines` interceptor. The (possibly modified) response returned by + `post_list_training_pipelines` will be passed to + `post_list_training_pipelines_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1057,6 +1297,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_cancel_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_cancel_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1218,6 +1465,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_delete_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_delete_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1623,6 +1877,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1786,6 +2044,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1942,6 +2207,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2096,6 +2365,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2251,6 +2527,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2408,6 +2688,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2565,6 +2849,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2720,6 +3008,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_training_pipelines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_training_pipelines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/prediction_service/async_client.py b/google/cloud/aiplatform_v1/services/prediction_service/async_client.py index e564489ecf..0283807411 100644 --- a/google/cloud/aiplatform_v1/services/prediction_service/async_client.py +++ b/google/cloud/aiplatform_v1/services/prediction_service/async_client.py @@ -58,6 +58,7 @@ from google.longrunning import operations_pb2 # type: ignore from google.protobuf import any_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from .transports.base import PredictionServiceTransport, DEFAULT_CLIENT_INFO from .transports.grpc_asyncio import PredictionServiceGrpcAsyncIOTransport from .client import PredictionServiceClient diff --git a/google/cloud/aiplatform_v1/services/prediction_service/client.py b/google/cloud/aiplatform_v1/services/prediction_service/client.py index f566757527..933ba75a28 100644 --- a/google/cloud/aiplatform_v1/services/prediction_service/client.py +++ b/google/cloud/aiplatform_v1/services/prediction_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -71,6 +73,7 @@ from google.longrunning import operations_pb2 # type: ignore from google.protobuf import any_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from .transports.base import PredictionServiceTransport, DEFAULT_CLIENT_INFO from .transports.grpc import PredictionServiceGrpcTransport from .transports.grpc_asyncio import PredictionServiceGrpcAsyncIOTransport @@ -574,6 +577,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2426,16 +2456,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2481,16 +2515,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2653,16 +2691,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2774,16 +2816,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2896,16 +2942,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2956,16 +3006,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3011,16 +3065,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3066,16 +3124,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/prediction_service/transports/rest.py b/google/cloud/aiplatform_v1/services/prediction_service/transports/rest.py index 2aee01cdf9..bb1a5088e7 100644 --- a/google/cloud/aiplatform_v1/services/prediction_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/prediction_service/transports/rest.py @@ -176,12 +176,38 @@ def post_direct_predict( ) -> prediction_service.DirectPredictResponse: """Post-rpc interceptor for direct_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_predict` interceptor runs + before the `post_direct_predict_with_metadata` interceptor. """ return response + def post_direct_predict_with_metadata( + self, + response: prediction_service.DirectPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_predict_with_metadata` + interceptor in new development instead of the `post_direct_predict` interceptor. + When both interceptors are used, this `post_direct_predict_with_metadata` interceptor runs after the + `post_direct_predict` interceptor. The (possibly modified) response returned by + `post_direct_predict` will be passed to + `post_direct_predict_with_metadata`. + """ + return response, metadata + def pre_direct_raw_predict( self, request: prediction_service.DirectRawPredictRequest, @@ -202,12 +228,38 @@ def post_direct_raw_predict( ) -> prediction_service.DirectRawPredictResponse: """Post-rpc interceptor for direct_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_raw_predict` interceptor runs + before the `post_direct_raw_predict_with_metadata` interceptor. """ return response + def post_direct_raw_predict_with_metadata( + self, + response: prediction_service.DirectRawPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectRawPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_raw_predict_with_metadata` + interceptor in new development instead of the `post_direct_raw_predict` interceptor. + When both interceptors are used, this `post_direct_raw_predict_with_metadata` interceptor runs after the + `post_direct_raw_predict` interceptor. The (possibly modified) response returned by + `post_direct_raw_predict` will be passed to + `post_direct_raw_predict_with_metadata`. + """ + return response, metadata + def pre_explain( self, request: prediction_service.ExplainRequest, @@ -227,12 +279,37 @@ def post_explain( ) -> prediction_service.ExplainResponse: """Post-rpc interceptor for explain - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_explain_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_explain` interceptor runs + before the `post_explain_with_metadata` interceptor. """ return response + def post_explain_with_metadata( + self, + response: prediction_service.ExplainResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.ExplainResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for explain + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_explain_with_metadata` + interceptor in new development instead of the `post_explain` interceptor. + When both interceptors are used, this `post_explain_with_metadata` interceptor runs after the + `post_explain` interceptor. The (possibly modified) response returned by + `post_explain` will be passed to + `post_explain_with_metadata`. + """ + return response, metadata + def pre_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -253,12 +330,38 @@ def post_generate_content( ) -> prediction_service.GenerateContentResponse: """Post-rpc interceptor for generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_generate_content` interceptor runs + before the `post_generate_content_with_metadata` interceptor. """ return response + def post_generate_content_with_metadata( + self, + response: prediction_service.GenerateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.GenerateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_generate_content_with_metadata` + interceptor in new development instead of the `post_generate_content` interceptor. + When both interceptors are used, this `post_generate_content_with_metadata` interceptor runs after the + `post_generate_content` interceptor. The (possibly modified) response returned by + `post_generate_content` will be passed to + `post_generate_content_with_metadata`. + """ + return response, metadata + def pre_predict( self, request: prediction_service.PredictRequest, @@ -278,12 +381,37 @@ def post_predict( ) -> prediction_service.PredictResponse: """Post-rpc interceptor for predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_predict` interceptor runs + before the `post_predict_with_metadata` interceptor. """ return response + def post_predict_with_metadata( + self, + response: prediction_service.PredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.PredictResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_predict_with_metadata` + interceptor in new development instead of the `post_predict` interceptor. + When both interceptors are used, this `post_predict_with_metadata` interceptor runs after the + `post_predict` interceptor. The (possibly modified) response returned by + `post_predict` will be passed to + `post_predict_with_metadata`. + """ + return response, metadata + def pre_raw_predict( self, request: prediction_service.RawPredictRequest, @@ -303,12 +431,35 @@ def post_raw_predict( ) -> httpbody_pb2.HttpBody: """Post-rpc interceptor for raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_raw_predict` interceptor runs + before the `post_raw_predict_with_metadata` interceptor. """ return response + def post_raw_predict_with_metadata( + self, + response: httpbody_pb2.HttpBody, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[httpbody_pb2.HttpBody, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_raw_predict_with_metadata` + interceptor in new development instead of the `post_raw_predict` interceptor. + When both interceptors are used, this `post_raw_predict_with_metadata` interceptor runs after the + `post_raw_predict` interceptor. The (possibly modified) response returned by + `post_raw_predict` will be passed to + `post_raw_predict_with_metadata`. + """ + return response, metadata + def pre_server_streaming_predict( self, request: prediction_service.StreamingPredictRequest, @@ -329,12 +480,37 @@ def post_server_streaming_predict( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for server_streaming_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_server_streaming_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_server_streaming_predict` interceptor runs + before the `post_server_streaming_predict_with_metadata` interceptor. """ return response + def post_server_streaming_predict_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for server_streaming_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_server_streaming_predict_with_metadata` + interceptor in new development instead of the `post_server_streaming_predict` interceptor. + When both interceptors are used, this `post_server_streaming_predict_with_metadata` interceptor runs after the + `post_server_streaming_predict` interceptor. The (possibly modified) response returned by + `post_server_streaming_predict` will be passed to + `post_server_streaming_predict_with_metadata`. + """ + return response, metadata + def pre_stream_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -355,12 +531,37 @@ def post_stream_generate_content( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_generate_content` interceptor runs + before the `post_stream_generate_content_with_metadata` interceptor. """ return response + def post_stream_generate_content_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_generate_content_with_metadata` + interceptor in new development instead of the `post_stream_generate_content` interceptor. + When both interceptors are used, this `post_stream_generate_content_with_metadata` interceptor runs after the + `post_stream_generate_content` interceptor. The (possibly modified) response returned by + `post_stream_generate_content` will be passed to + `post_stream_generate_content_with_metadata`. + """ + return response, metadata + def pre_stream_raw_predict( self, request: prediction_service.StreamRawPredictRequest, @@ -381,12 +582,37 @@ def post_stream_raw_predict( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_raw_predict` interceptor runs + before the `post_stream_raw_predict_with_metadata` interceptor. """ return response + def post_stream_raw_predict_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_raw_predict_with_metadata` + interceptor in new development instead of the `post_stream_raw_predict` interceptor. + When both interceptors are used, this `post_stream_raw_predict_with_metadata` interceptor runs after the + `post_stream_raw_predict` interceptor. The (possibly modified) response returned by + `post_stream_raw_predict` will be passed to + `post_stream_raw_predict_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -845,6 +1071,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_direct_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_direct_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1000,6 +1230,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_direct_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_direct_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1156,6 +1390,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_explain(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_explain_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1310,6 +1548,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_generate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1466,6 +1708,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1665,6 +1911,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1825,6 +2075,10 @@ def __call__( ) resp = self._interceptor.post_server_streaming_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_server_streaming_predict_with_metadata( + resp, response_metadata + ) return resp class _StreamDirectPredict( @@ -1996,6 +2250,10 @@ def __call__( ) resp = self._interceptor.post_stream_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stream_generate_content_with_metadata( + resp, response_metadata + ) return resp class _StreamingPredict( @@ -2211,6 +2469,10 @@ def __call__( resp = rest_streaming.ResponseIterator(response, httpbody_pb2.HttpBody) resp = self._interceptor.post_stream_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stream_raw_predict_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1/services/prediction_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/prediction_service/transports/rest_asyncio.py index aceded8522..d25dd02e71 100644 --- a/google/cloud/aiplatform_v1/services/prediction_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/prediction_service/transports/rest_asyncio.py @@ -193,12 +193,38 @@ async def post_direct_predict( ) -> prediction_service.DirectPredictResponse: """Post-rpc interceptor for direct_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_predict` interceptor runs + before the `post_direct_predict_with_metadata` interceptor. """ return response + async def post_direct_predict_with_metadata( + self, + response: prediction_service.DirectPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_predict_with_metadata` + interceptor in new development instead of the `post_direct_predict` interceptor. + When both interceptors are used, this `post_direct_predict_with_metadata` interceptor runs after the + `post_direct_predict` interceptor. The (possibly modified) response returned by + `post_direct_predict` will be passed to + `post_direct_predict_with_metadata`. + """ + return response, metadata + async def pre_direct_raw_predict( self, request: prediction_service.DirectRawPredictRequest, @@ -219,12 +245,38 @@ async def post_direct_raw_predict( ) -> prediction_service.DirectRawPredictResponse: """Post-rpc interceptor for direct_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_raw_predict` interceptor runs + before the `post_direct_raw_predict_with_metadata` interceptor. """ return response + async def post_direct_raw_predict_with_metadata( + self, + response: prediction_service.DirectRawPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectRawPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_raw_predict_with_metadata` + interceptor in new development instead of the `post_direct_raw_predict` interceptor. + When both interceptors are used, this `post_direct_raw_predict_with_metadata` interceptor runs after the + `post_direct_raw_predict` interceptor. The (possibly modified) response returned by + `post_direct_raw_predict` will be passed to + `post_direct_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_explain( self, request: prediction_service.ExplainRequest, @@ -244,12 +296,37 @@ async def post_explain( ) -> prediction_service.ExplainResponse: """Post-rpc interceptor for explain - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_explain_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_explain` interceptor runs + before the `post_explain_with_metadata` interceptor. """ return response + async def post_explain_with_metadata( + self, + response: prediction_service.ExplainResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.ExplainResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for explain + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_explain_with_metadata` + interceptor in new development instead of the `post_explain` interceptor. + When both interceptors are used, this `post_explain_with_metadata` interceptor runs after the + `post_explain` interceptor. The (possibly modified) response returned by + `post_explain` will be passed to + `post_explain_with_metadata`. + """ + return response, metadata + async def pre_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -270,12 +347,38 @@ async def post_generate_content( ) -> prediction_service.GenerateContentResponse: """Post-rpc interceptor for generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_generate_content` interceptor runs + before the `post_generate_content_with_metadata` interceptor. """ return response + async def post_generate_content_with_metadata( + self, + response: prediction_service.GenerateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.GenerateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_generate_content_with_metadata` + interceptor in new development instead of the `post_generate_content` interceptor. + When both interceptors are used, this `post_generate_content_with_metadata` interceptor runs after the + `post_generate_content` interceptor. The (possibly modified) response returned by + `post_generate_content` will be passed to + `post_generate_content_with_metadata`. + """ + return response, metadata + async def pre_predict( self, request: prediction_service.PredictRequest, @@ -295,12 +398,37 @@ async def post_predict( ) -> prediction_service.PredictResponse: """Post-rpc interceptor for predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_predict` interceptor runs + before the `post_predict_with_metadata` interceptor. """ return response + async def post_predict_with_metadata( + self, + response: prediction_service.PredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.PredictResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_predict_with_metadata` + interceptor in new development instead of the `post_predict` interceptor. + When both interceptors are used, this `post_predict_with_metadata` interceptor runs after the + `post_predict` interceptor. The (possibly modified) response returned by + `post_predict` will be passed to + `post_predict_with_metadata`. + """ + return response, metadata + async def pre_raw_predict( self, request: prediction_service.RawPredictRequest, @@ -320,12 +448,35 @@ async def post_raw_predict( ) -> httpbody_pb2.HttpBody: """Post-rpc interceptor for raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_raw_predict` interceptor runs + before the `post_raw_predict_with_metadata` interceptor. """ return response + async def post_raw_predict_with_metadata( + self, + response: httpbody_pb2.HttpBody, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[httpbody_pb2.HttpBody, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_raw_predict_with_metadata` + interceptor in new development instead of the `post_raw_predict` interceptor. + When both interceptors are used, this `post_raw_predict_with_metadata` interceptor runs after the + `post_raw_predict` interceptor. The (possibly modified) response returned by + `post_raw_predict` will be passed to + `post_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_server_streaming_predict( self, request: prediction_service.StreamingPredictRequest, @@ -346,12 +497,38 @@ async def post_server_streaming_predict( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for server_streaming_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_server_streaming_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_server_streaming_predict` interceptor runs + before the `post_server_streaming_predict_with_metadata` interceptor. """ return response + async def post_server_streaming_predict_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for server_streaming_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_server_streaming_predict_with_metadata` + interceptor in new development instead of the `post_server_streaming_predict` interceptor. + When both interceptors are used, this `post_server_streaming_predict_with_metadata` interceptor runs after the + `post_server_streaming_predict` interceptor. The (possibly modified) response returned by + `post_server_streaming_predict` will be passed to + `post_server_streaming_predict_with_metadata`. + """ + return response, metadata + async def pre_stream_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -372,12 +549,38 @@ async def post_stream_generate_content( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_generate_content` interceptor runs + before the `post_stream_generate_content_with_metadata` interceptor. """ return response + async def post_stream_generate_content_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_generate_content_with_metadata` + interceptor in new development instead of the `post_stream_generate_content` interceptor. + When both interceptors are used, this `post_stream_generate_content_with_metadata` interceptor runs after the + `post_stream_generate_content` interceptor. The (possibly modified) response returned by + `post_stream_generate_content` will be passed to + `post_stream_generate_content_with_metadata`. + """ + return response, metadata + async def pre_stream_raw_predict( self, request: prediction_service.StreamRawPredictRequest, @@ -398,12 +601,38 @@ async def post_stream_raw_predict( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_raw_predict` interceptor runs + before the `post_stream_raw_predict_with_metadata` interceptor. """ return response + async def post_stream_raw_predict_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_raw_predict_with_metadata` + interceptor in new development instead of the `post_stream_raw_predict` interceptor. + When both interceptors are used, this `post_stream_raw_predict_with_metadata` interceptor runs after the + `post_stream_raw_predict` interceptor. The (possibly modified) response returned by + `post_stream_raw_predict` will be passed to + `post_stream_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -975,6 +1204,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_direct_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_direct_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1136,6 +1369,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_direct_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_direct_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1298,6 +1535,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_explain(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_explain_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1458,6 +1699,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_generate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1620,6 +1865,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1830,6 +2079,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1993,6 +2246,13 @@ async def __call__( response, prediction_service.StreamingPredictResponse ) resp = await self._interceptor.post_server_streaming_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_server_streaming_predict_with_metadata( + resp, response_metadata + ) return resp class _StreamDirectPredict( @@ -2166,6 +2426,13 @@ async def __call__( response, prediction_service.GenerateContentResponse ) resp = await self._interceptor.post_stream_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_stream_generate_content_with_metadata( + resp, response_metadata + ) return resp class _StreamingPredict( @@ -2387,6 +2654,10 @@ async def __call__( response, httpbody_pb2.HttpBody ) resp = await self._interceptor.post_stream_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stream_raw_predict_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py index bf5a579975..f463dd914b 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -517,6 +519,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1041,16 +1070,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1096,16 +1129,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1268,16 +1305,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1389,16 +1430,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1511,16 +1556,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1571,16 +1620,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1626,16 +1679,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1681,16 +1738,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest.py b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest.py index cfb79f5022..b1ff335ef8 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest.py @@ -121,12 +121,38 @@ def post_query_reasoning_engine( ) -> reasoning_engine_execution_service.QueryReasoningEngineResponse: """Post-rpc interceptor for query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_reasoning_engine` interceptor runs + before the `post_query_reasoning_engine_with_metadata` interceptor. """ return response + def post_query_reasoning_engine_with_metadata( + self, + response: reasoning_engine_execution_service.QueryReasoningEngineResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_execution_service.QueryReasoningEngineResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_query_reasoning_engine_with_metadata` interceptor runs after the + `post_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_query_reasoning_engine` will be passed to + `post_query_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_stream_query_reasoning_engine( self, request: reasoning_engine_execution_service.StreamQueryReasoningEngineRequest, @@ -147,12 +173,37 @@ def post_stream_query_reasoning_engine( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_query_reasoning_engine` interceptor runs + before the `post_stream_query_reasoning_engine_with_metadata` interceptor. """ return response + def post_stream_query_reasoning_engine_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_stream_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_stream_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_stream_query_reasoning_engine_with_metadata` interceptor runs after the + `post_stream_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_stream_query_reasoning_engine` will be passed to + `post_stream_query_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -621,6 +672,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_query_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -824,6 +879,13 @@ def __call__( resp = rest_streaming.ResponseIterator(response, httpbody_pb2.HttpBody) resp = self._interceptor.post_stream_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_stream_query_reasoning_engine_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest_asyncio.py index cd039908ec..a8de339d82 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/transports/rest_asyncio.py @@ -138,12 +138,38 @@ async def post_query_reasoning_engine( ) -> reasoning_engine_execution_service.QueryReasoningEngineResponse: """Post-rpc interceptor for query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_reasoning_engine` interceptor runs + before the `post_query_reasoning_engine_with_metadata` interceptor. """ return response + async def post_query_reasoning_engine_with_metadata( + self, + response: reasoning_engine_execution_service.QueryReasoningEngineResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_execution_service.QueryReasoningEngineResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_query_reasoning_engine_with_metadata` interceptor runs after the + `post_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_query_reasoning_engine` will be passed to + `post_query_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_stream_query_reasoning_engine( self, request: reasoning_engine_execution_service.StreamQueryReasoningEngineRequest, @@ -164,12 +190,38 @@ async def post_stream_query_reasoning_engine( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_query_reasoning_engine` interceptor runs + before the `post_stream_query_reasoning_engine_with_metadata` interceptor. """ return response + async def post_stream_query_reasoning_engine_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_stream_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_stream_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_stream_query_reasoning_engine_with_metadata` interceptor runs after the + `post_stream_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_stream_query_reasoning_engine` will be passed to + `post_stream_query_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -694,6 +746,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_query_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -907,6 +963,13 @@ async def __call__( response, httpbody_pb2.HttpBody ) resp = await self._interceptor.post_stream_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_stream_query_reasoning_engine_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_service/client.py b/google/cloud/aiplatform_v1/services/reasoning_engine_service/client.py index 66f336b4da..889932382a 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_service/client.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -511,6 +513,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1437,16 +1466,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1492,16 +1525,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1664,16 +1701,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1785,16 +1826,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1907,16 +1952,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1967,16 +2016,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2022,16 +2075,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2077,16 +2134,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest.py b/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest.py index ae6756ac27..2a3c2ea21c 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest.py @@ -146,12 +146,35 @@ def post_create_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_create_reasoning_engine` interceptor runs + before the `post_create_reasoning_engine_with_metadata` interceptor. """ return response + def post_create_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_create_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_create_reasoning_engine` interceptor. + When both interceptors are used, this `post_create_reasoning_engine_with_metadata` interceptor runs after the + `post_create_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_create_reasoning_engine` will be passed to + `post_create_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_delete_reasoning_engine( self, request: reasoning_engine_service.DeleteReasoningEngineRequest, @@ -172,12 +195,35 @@ def post_delete_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_reasoning_engine` interceptor runs + before the `post_delete_reasoning_engine_with_metadata` interceptor. """ return response + def post_delete_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_delete_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_delete_reasoning_engine` interceptor. + When both interceptors are used, this `post_delete_reasoning_engine_with_metadata` interceptor runs after the + `post_delete_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_delete_reasoning_engine` will be passed to + `post_delete_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_reasoning_engine( self, request: reasoning_engine_service.GetReasoningEngineRequest, @@ -198,12 +244,37 @@ def post_get_reasoning_engine( ) -> reasoning_engine.ReasoningEngine: """Post-rpc interceptor for get_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_get_reasoning_engine` interceptor runs + before the `post_get_reasoning_engine_with_metadata` interceptor. """ return response + def post_get_reasoning_engine_with_metadata( + self, + response: reasoning_engine.ReasoningEngine, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine.ReasoningEngine, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_get_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_get_reasoning_engine` interceptor. + When both interceptors are used, this `post_get_reasoning_engine_with_metadata` interceptor runs after the + `post_get_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_get_reasoning_engine` will be passed to + `post_get_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_list_reasoning_engines( self, request: reasoning_engine_service.ListReasoningEnginesRequest, @@ -224,12 +295,38 @@ def post_list_reasoning_engines( ) -> reasoning_engine_service.ListReasoningEnginesResponse: """Post-rpc interceptor for list_reasoning_engines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_reasoning_engines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_list_reasoning_engines` interceptor runs + before the `post_list_reasoning_engines_with_metadata` interceptor. """ return response + def post_list_reasoning_engines_with_metadata( + self, + response: reasoning_engine_service.ListReasoningEnginesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_service.ListReasoningEnginesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_reasoning_engines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_list_reasoning_engines_with_metadata` + interceptor in new development instead of the `post_list_reasoning_engines` interceptor. + When both interceptors are used, this `post_list_reasoning_engines_with_metadata` interceptor runs after the + `post_list_reasoning_engines` interceptor. The (possibly modified) response returned by + `post_list_reasoning_engines` will be passed to + `post_list_reasoning_engines_with_metadata`. + """ + return response, metadata + def pre_update_reasoning_engine( self, request: reasoning_engine_service.UpdateReasoningEngineRequest, @@ -250,12 +347,35 @@ def post_update_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_update_reasoning_engine` interceptor runs + before the `post_update_reasoning_engine_with_metadata` interceptor. """ return response + def post_update_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_update_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_update_reasoning_engine` interceptor. + When both interceptors are used, this `post_update_reasoning_engine_with_metadata` interceptor runs after the + `post_update_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_update_reasoning_engine` will be passed to + `post_update_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2630,6 +2750,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2776,6 +2900,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2927,6 +3055,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3078,6 +3210,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_reasoning_engines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_reasoning_engines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3234,6 +3370,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest_asyncio.py index 4641248c10..977a4196d2 100644 --- a/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/reasoning_engine_service/transports/rest_asyncio.py @@ -164,12 +164,35 @@ async def post_create_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_create_reasoning_engine` interceptor runs + before the `post_create_reasoning_engine_with_metadata` interceptor. """ return response + async def post_create_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_create_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_create_reasoning_engine` interceptor. + When both interceptors are used, this `post_create_reasoning_engine_with_metadata` interceptor runs after the + `post_create_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_create_reasoning_engine` will be passed to + `post_create_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_delete_reasoning_engine( self, request: reasoning_engine_service.DeleteReasoningEngineRequest, @@ -190,12 +213,35 @@ async def post_delete_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_reasoning_engine` interceptor runs + before the `post_delete_reasoning_engine_with_metadata` interceptor. """ return response + async def post_delete_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_delete_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_delete_reasoning_engine` interceptor. + When both interceptors are used, this `post_delete_reasoning_engine_with_metadata` interceptor runs after the + `post_delete_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_delete_reasoning_engine` will be passed to + `post_delete_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_reasoning_engine( self, request: reasoning_engine_service.GetReasoningEngineRequest, @@ -216,12 +262,37 @@ async def post_get_reasoning_engine( ) -> reasoning_engine.ReasoningEngine: """Post-rpc interceptor for get_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_get_reasoning_engine` interceptor runs + before the `post_get_reasoning_engine_with_metadata` interceptor. """ return response + async def post_get_reasoning_engine_with_metadata( + self, + response: reasoning_engine.ReasoningEngine, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine.ReasoningEngine, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_get_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_get_reasoning_engine` interceptor. + When both interceptors are used, this `post_get_reasoning_engine_with_metadata` interceptor runs after the + `post_get_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_get_reasoning_engine` will be passed to + `post_get_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_list_reasoning_engines( self, request: reasoning_engine_service.ListReasoningEnginesRequest, @@ -242,12 +313,38 @@ async def post_list_reasoning_engines( ) -> reasoning_engine_service.ListReasoningEnginesResponse: """Post-rpc interceptor for list_reasoning_engines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_reasoning_engines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_list_reasoning_engines` interceptor runs + before the `post_list_reasoning_engines_with_metadata` interceptor. """ return response + async def post_list_reasoning_engines_with_metadata( + self, + response: reasoning_engine_service.ListReasoningEnginesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_service.ListReasoningEnginesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_reasoning_engines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_list_reasoning_engines_with_metadata` + interceptor in new development instead of the `post_list_reasoning_engines` interceptor. + When both interceptors are used, this `post_list_reasoning_engines_with_metadata` interceptor runs after the + `post_list_reasoning_engines` interceptor. The (possibly modified) response returned by + `post_list_reasoning_engines` will be passed to + `post_list_reasoning_engines_with_metadata`. + """ + return response, metadata + async def pre_update_reasoning_engine( self, request: reasoning_engine_service.UpdateReasoningEngineRequest, @@ -268,12 +365,35 @@ async def post_update_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_update_reasoning_engine` interceptor runs + before the `post_update_reasoning_engine_with_metadata` interceptor. """ return response + async def post_update_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_update_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_update_reasoning_engine` interceptor. + When both interceptors are used, this `post_update_reasoning_engine_with_metadata` interceptor runs after the + `post_update_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_update_reasoning_engine` will be passed to + `post_update_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -811,6 +931,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -967,6 +1094,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1122,6 +1256,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1277,6 +1415,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_reasoning_engines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_reasoning_engines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1443,6 +1585,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/schedule_service/client.py b/google/cloud/aiplatform_v1/services/schedule_service/client.py index cbf5650550..f5f579adeb 100644 --- a/google/cloud/aiplatform_v1/services/schedule_service/client.py +++ b/google/cloud/aiplatform_v1/services/schedule_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -737,6 +739,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1846,16 +1875,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1901,16 +1934,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2073,16 +2110,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2194,16 +2235,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2316,16 +2361,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2376,16 +2425,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2431,16 +2484,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2486,16 +2543,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/schedule_service/transports/rest.py b/google/cloud/aiplatform_v1/services/schedule_service/transports/rest.py index 0594e76aed..8d99df163c 100644 --- a/google/cloud/aiplatform_v1/services/schedule_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/schedule_service/transports/rest.py @@ -155,12 +155,35 @@ def post_create_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for create_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_create_schedule` interceptor runs + before the `post_create_schedule_with_metadata` interceptor. """ return response + def post_create_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_create_schedule_with_metadata` + interceptor in new development instead of the `post_create_schedule` interceptor. + When both interceptors are used, this `post_create_schedule_with_metadata` interceptor runs after the + `post_create_schedule` interceptor. The (possibly modified) response returned by + `post_create_schedule` will be passed to + `post_create_schedule_with_metadata`. + """ + return response, metadata + def pre_delete_schedule( self, request: schedule_service.DeleteScheduleRequest, @@ -180,12 +203,35 @@ def post_delete_schedule( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_delete_schedule` interceptor runs + before the `post_delete_schedule_with_metadata` interceptor. """ return response + def post_delete_schedule_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_delete_schedule_with_metadata` + interceptor in new development instead of the `post_delete_schedule` interceptor. + When both interceptors are used, this `post_delete_schedule_with_metadata` interceptor runs after the + `post_delete_schedule` interceptor. The (possibly modified) response returned by + `post_delete_schedule` will be passed to + `post_delete_schedule_with_metadata`. + """ + return response, metadata + def pre_get_schedule( self, request: schedule_service.GetScheduleRequest, @@ -203,12 +249,35 @@ def pre_get_schedule( def post_get_schedule(self, response: schedule.Schedule) -> schedule.Schedule: """Post-rpc interceptor for get_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_get_schedule` interceptor runs + before the `post_get_schedule_with_metadata` interceptor. """ return response + def post_get_schedule_with_metadata( + self, + response: schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_get_schedule_with_metadata` + interceptor in new development instead of the `post_get_schedule` interceptor. + When both interceptors are used, this `post_get_schedule_with_metadata` interceptor runs after the + `post_get_schedule` interceptor. The (possibly modified) response returned by + `post_get_schedule` will be passed to + `post_get_schedule_with_metadata`. + """ + return response, metadata + def pre_list_schedules( self, request: schedule_service.ListSchedulesRequest, @@ -228,12 +297,37 @@ def post_list_schedules( ) -> schedule_service.ListSchedulesResponse: """Post-rpc interceptor for list_schedules - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_schedules_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_list_schedules` interceptor runs + before the `post_list_schedules_with_metadata` interceptor. """ return response + def post_list_schedules_with_metadata( + self, + response: schedule_service.ListSchedulesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + schedule_service.ListSchedulesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_schedules + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_list_schedules_with_metadata` + interceptor in new development instead of the `post_list_schedules` interceptor. + When both interceptors are used, this `post_list_schedules_with_metadata` interceptor runs after the + `post_list_schedules` interceptor. The (possibly modified) response returned by + `post_list_schedules` will be passed to + `post_list_schedules_with_metadata`. + """ + return response, metadata + def pre_pause_schedule( self, request: schedule_service.PauseScheduleRequest, @@ -281,12 +375,35 @@ def post_update_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for update_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_update_schedule` interceptor runs + before the `post_update_schedule_with_metadata` interceptor. """ return response + def post_update_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_update_schedule_with_metadata` + interceptor in new development instead of the `post_update_schedule` interceptor. + When both interceptors are used, this `post_update_schedule_with_metadata` interceptor runs after the + `post_update_schedule` interceptor. The (possibly modified) response returned by + `post_update_schedule` will be passed to + `post_update_schedule_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2663,6 +2780,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2806,6 +2927,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2952,6 +3077,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3096,6 +3225,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_schedules(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_schedules_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3478,6 +3611,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/schedule_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/schedule_service/transports/rest_asyncio.py index 7556dcf78a..d2929d3cbc 100644 --- a/google/cloud/aiplatform_v1/services/schedule_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/schedule_service/transports/rest_asyncio.py @@ -173,12 +173,35 @@ async def post_create_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for create_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_create_schedule` interceptor runs + before the `post_create_schedule_with_metadata` interceptor. """ return response + async def post_create_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_create_schedule_with_metadata` + interceptor in new development instead of the `post_create_schedule` interceptor. + When both interceptors are used, this `post_create_schedule_with_metadata` interceptor runs after the + `post_create_schedule` interceptor. The (possibly modified) response returned by + `post_create_schedule` will be passed to + `post_create_schedule_with_metadata`. + """ + return response, metadata + async def pre_delete_schedule( self, request: schedule_service.DeleteScheduleRequest, @@ -198,12 +221,35 @@ async def post_delete_schedule( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_delete_schedule` interceptor runs + before the `post_delete_schedule_with_metadata` interceptor. """ return response + async def post_delete_schedule_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_delete_schedule_with_metadata` + interceptor in new development instead of the `post_delete_schedule` interceptor. + When both interceptors are used, this `post_delete_schedule_with_metadata` interceptor runs after the + `post_delete_schedule` interceptor. The (possibly modified) response returned by + `post_delete_schedule` will be passed to + `post_delete_schedule_with_metadata`. + """ + return response, metadata + async def pre_get_schedule( self, request: schedule_service.GetScheduleRequest, @@ -221,12 +267,35 @@ async def pre_get_schedule( async def post_get_schedule(self, response: schedule.Schedule) -> schedule.Schedule: """Post-rpc interceptor for get_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_get_schedule` interceptor runs + before the `post_get_schedule_with_metadata` interceptor. """ return response + async def post_get_schedule_with_metadata( + self, + response: schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_get_schedule_with_metadata` + interceptor in new development instead of the `post_get_schedule` interceptor. + When both interceptors are used, this `post_get_schedule_with_metadata` interceptor runs after the + `post_get_schedule` interceptor. The (possibly modified) response returned by + `post_get_schedule` will be passed to + `post_get_schedule_with_metadata`. + """ + return response, metadata + async def pre_list_schedules( self, request: schedule_service.ListSchedulesRequest, @@ -246,12 +315,37 @@ async def post_list_schedules( ) -> schedule_service.ListSchedulesResponse: """Post-rpc interceptor for list_schedules - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_schedules_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_list_schedules` interceptor runs + before the `post_list_schedules_with_metadata` interceptor. """ return response + async def post_list_schedules_with_metadata( + self, + response: schedule_service.ListSchedulesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + schedule_service.ListSchedulesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_schedules + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_list_schedules_with_metadata` + interceptor in new development instead of the `post_list_schedules` interceptor. + When both interceptors are used, this `post_list_schedules_with_metadata` interceptor runs after the + `post_list_schedules` interceptor. The (possibly modified) response returned by + `post_list_schedules` will be passed to + `post_list_schedules_with_metadata`. + """ + return response, metadata + async def pre_pause_schedule( self, request: schedule_service.PauseScheduleRequest, @@ -299,12 +393,35 @@ async def post_update_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for update_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_update_schedule` interceptor runs + before the `post_update_schedule_with_metadata` interceptor. """ return response + async def post_update_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_update_schedule_with_metadata` + interceptor in new development instead of the `post_update_schedule` interceptor. + When both interceptors are used, this `post_update_schedule_with_metadata` interceptor runs after the + `post_update_schedule` interceptor. The (possibly modified) response returned by + `post_update_schedule` will be passed to + `post_update_schedule_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -853,6 +970,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1009,6 +1130,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1165,6 +1290,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1320,6 +1449,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_schedules(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_schedules_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1735,6 +1868,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/specialist_pool_service/client.py b/google/cloud/aiplatform_v1/services/specialist_pool_service/client.py index 765334c6ed..e6e82a5908 100644 --- a/google/cloud/aiplatform_v1/services/specialist_pool_service/client.py +++ b/google/cloud/aiplatform_v1/services/specialist_pool_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -516,6 +518,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1457,16 +1486,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1512,16 +1545,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1684,16 +1721,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1805,16 +1846,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1927,16 +1972,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1987,16 +2036,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2042,16 +2095,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2097,16 +2154,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest.py b/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest.py index 356fe2c9d8..0df5712d13 100644 --- a/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest.py @@ -146,12 +146,35 @@ def post_create_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_specialist_pool` interceptor runs + before the `post_create_specialist_pool_with_metadata` interceptor. """ return response + def post_create_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_create_specialist_pool_with_metadata` + interceptor in new development instead of the `post_create_specialist_pool` interceptor. + When both interceptors are used, this `post_create_specialist_pool_with_metadata` interceptor runs after the + `post_create_specialist_pool` interceptor. The (possibly modified) response returned by + `post_create_specialist_pool` will be passed to + `post_create_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_delete_specialist_pool( self, request: specialist_pool_service.DeleteSpecialistPoolRequest, @@ -172,12 +195,35 @@ def post_delete_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_specialist_pool` interceptor runs + before the `post_delete_specialist_pool_with_metadata` interceptor. """ return response + def post_delete_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_delete_specialist_pool_with_metadata` + interceptor in new development instead of the `post_delete_specialist_pool` interceptor. + When both interceptors are used, this `post_delete_specialist_pool_with_metadata` interceptor runs after the + `post_delete_specialist_pool` interceptor. The (possibly modified) response returned by + `post_delete_specialist_pool` will be passed to + `post_delete_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_get_specialist_pool( self, request: specialist_pool_service.GetSpecialistPoolRequest, @@ -198,12 +244,35 @@ def post_get_specialist_pool( ) -> specialist_pool.SpecialistPool: """Post-rpc interceptor for get_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_specialist_pool` interceptor runs + before the `post_get_specialist_pool_with_metadata` interceptor. """ return response + def post_get_specialist_pool_with_metadata( + self, + response: specialist_pool.SpecialistPool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[specialist_pool.SpecialistPool, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_get_specialist_pool_with_metadata` + interceptor in new development instead of the `post_get_specialist_pool` interceptor. + When both interceptors are used, this `post_get_specialist_pool_with_metadata` interceptor runs after the + `post_get_specialist_pool` interceptor. The (possibly modified) response returned by + `post_get_specialist_pool` will be passed to + `post_get_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_list_specialist_pools( self, request: specialist_pool_service.ListSpecialistPoolsRequest, @@ -224,12 +293,38 @@ def post_list_specialist_pools( ) -> specialist_pool_service.ListSpecialistPoolsResponse: """Post-rpc interceptor for list_specialist_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_specialist_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_specialist_pools` interceptor runs + before the `post_list_specialist_pools_with_metadata` interceptor. """ return response + def post_list_specialist_pools_with_metadata( + self, + response: specialist_pool_service.ListSpecialistPoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + specialist_pool_service.ListSpecialistPoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_specialist_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_list_specialist_pools_with_metadata` + interceptor in new development instead of the `post_list_specialist_pools` interceptor. + When both interceptors are used, this `post_list_specialist_pools_with_metadata` interceptor runs after the + `post_list_specialist_pools` interceptor. The (possibly modified) response returned by + `post_list_specialist_pools` will be passed to + `post_list_specialist_pools_with_metadata`. + """ + return response, metadata + def pre_update_specialist_pool( self, request: specialist_pool_service.UpdateSpecialistPoolRequest, @@ -250,12 +345,35 @@ def post_update_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_specialist_pool` interceptor runs + before the `post_update_specialist_pool_with_metadata` interceptor. """ return response + def post_update_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_update_specialist_pool_with_metadata` + interceptor in new development instead of the `post_update_specialist_pool` interceptor. + When both interceptors are used, this `post_update_specialist_pool_with_metadata` interceptor runs after the + `post_update_specialist_pool` interceptor. The (possibly modified) response returned by + `post_update_specialist_pool` will be passed to + `post_update_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2637,6 +2755,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2785,6 +2907,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2944,6 +3070,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3093,6 +3223,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_specialist_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_specialist_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3251,6 +3385,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest_asyncio.py index 78a118ab28..4d55083607 100644 --- a/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/specialist_pool_service/transports/rest_asyncio.py @@ -164,12 +164,35 @@ async def post_create_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_specialist_pool` interceptor runs + before the `post_create_specialist_pool_with_metadata` interceptor. """ return response + async def post_create_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_create_specialist_pool_with_metadata` + interceptor in new development instead of the `post_create_specialist_pool` interceptor. + When both interceptors are used, this `post_create_specialist_pool_with_metadata` interceptor runs after the + `post_create_specialist_pool` interceptor. The (possibly modified) response returned by + `post_create_specialist_pool` will be passed to + `post_create_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_delete_specialist_pool( self, request: specialist_pool_service.DeleteSpecialistPoolRequest, @@ -190,12 +213,35 @@ async def post_delete_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_specialist_pool` interceptor runs + before the `post_delete_specialist_pool_with_metadata` interceptor. """ return response + async def post_delete_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_delete_specialist_pool_with_metadata` + interceptor in new development instead of the `post_delete_specialist_pool` interceptor. + When both interceptors are used, this `post_delete_specialist_pool_with_metadata` interceptor runs after the + `post_delete_specialist_pool` interceptor. The (possibly modified) response returned by + `post_delete_specialist_pool` will be passed to + `post_delete_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_get_specialist_pool( self, request: specialist_pool_service.GetSpecialistPoolRequest, @@ -216,12 +262,35 @@ async def post_get_specialist_pool( ) -> specialist_pool.SpecialistPool: """Post-rpc interceptor for get_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_specialist_pool` interceptor runs + before the `post_get_specialist_pool_with_metadata` interceptor. """ return response + async def post_get_specialist_pool_with_metadata( + self, + response: specialist_pool.SpecialistPool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[specialist_pool.SpecialistPool, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_get_specialist_pool_with_metadata` + interceptor in new development instead of the `post_get_specialist_pool` interceptor. + When both interceptors are used, this `post_get_specialist_pool_with_metadata` interceptor runs after the + `post_get_specialist_pool` interceptor. The (possibly modified) response returned by + `post_get_specialist_pool` will be passed to + `post_get_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_list_specialist_pools( self, request: specialist_pool_service.ListSpecialistPoolsRequest, @@ -242,12 +311,38 @@ async def post_list_specialist_pools( ) -> specialist_pool_service.ListSpecialistPoolsResponse: """Post-rpc interceptor for list_specialist_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_specialist_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_specialist_pools` interceptor runs + before the `post_list_specialist_pools_with_metadata` interceptor. """ return response + async def post_list_specialist_pools_with_metadata( + self, + response: specialist_pool_service.ListSpecialistPoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + specialist_pool_service.ListSpecialistPoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_specialist_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_list_specialist_pools_with_metadata` + interceptor in new development instead of the `post_list_specialist_pools` interceptor. + When both interceptors are used, this `post_list_specialist_pools_with_metadata` interceptor runs after the + `post_list_specialist_pools` interceptor. The (possibly modified) response returned by + `post_list_specialist_pools` will be passed to + `post_list_specialist_pools_with_metadata`. + """ + return response, metadata + async def pre_update_specialist_pool( self, request: specialist_pool_service.UpdateSpecialistPoolRequest, @@ -268,12 +363,35 @@ async def post_update_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_specialist_pool` interceptor runs + before the `post_update_specialist_pool_with_metadata` interceptor. """ return response + async def post_update_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_update_specialist_pool_with_metadata` + interceptor in new development instead of the `post_update_specialist_pool` interceptor. + When both interceptors are used, this `post_update_specialist_pool_with_metadata` interceptor runs after the + `post_update_specialist_pool` interceptor. The (possibly modified) response returned by + `post_update_specialist_pool` will be passed to + `post_update_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -812,6 +930,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -966,6 +1088,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1129,6 +1255,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1282,6 +1412,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_specialist_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_specialist_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1446,6 +1580,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/tensorboard_service/client.py b/google/cloud/aiplatform_v1/services/tensorboard_service/client.py index f6a7ff5e06..1e17804f14 100644 --- a/google/cloud/aiplatform_v1/services/tensorboard_service/client.py +++ b/google/cloud/aiplatform_v1/services/tensorboard_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -602,6 +604,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -4706,16 +4735,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -4761,16 +4794,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -4933,16 +4970,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5054,16 +5095,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5176,16 +5221,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5236,16 +5285,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -5291,16 +5344,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -5346,16 +5403,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest.py b/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest.py index 05b9c5e496..ffe120be4e 100644 --- a/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest.py @@ -356,12 +356,38 @@ def post_batch_create_tensorboard_runs( ) -> tensorboard_service.BatchCreateTensorboardRunsResponse: """Post-rpc interceptor for batch_create_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_runs` interceptor runs + before the `post_batch_create_tensorboard_runs_with_metadata` interceptor. """ return response + def post_batch_create_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_runs` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_runs_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_runs` will be passed to + `post_batch_create_tensorboard_runs_with_metadata`. + """ + return response, metadata + def pre_batch_create_tensorboard_time_series( self, request: tensorboard_service.BatchCreateTensorboardTimeSeriesRequest, @@ -382,12 +408,38 @@ def post_batch_create_tensorboard_time_series( ) -> tensorboard_service.BatchCreateTensorboardTimeSeriesResponse: """Post-rpc interceptor for batch_create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_time_series` interceptor runs + before the `post_batch_create_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_batch_create_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_time_series` will be passed to + `post_batch_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_batch_read_tensorboard_time_series_data( self, request: tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest, @@ -408,12 +460,38 @@ def post_batch_read_tensorboard_time_series_data( ) -> tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for batch_read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_tensorboard_time_series_data` interceptor runs + before the `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_batch_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_batch_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_batch_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_batch_read_tensorboard_time_series_data` will be passed to + `post_batch_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard( self, request: tensorboard_service.CreateTensorboardRequest, @@ -434,12 +512,35 @@ def post_create_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard` interceptor runs + before the `post_create_tensorboard_with_metadata` interceptor. """ return response + def post_create_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_with_metadata` + interceptor in new development instead of the `post_create_tensorboard` interceptor. + When both interceptors are used, this `post_create_tensorboard_with_metadata` interceptor runs after the + `post_create_tensorboard` interceptor. The (possibly modified) response returned by + `post_create_tensorboard` will be passed to + `post_create_tensorboard_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_experiment( self, request: tensorboard_service.CreateTensorboardExperimentRequest, @@ -460,12 +561,38 @@ def post_create_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for create_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_experiment` interceptor runs + before the `post_create_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_create_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_create_tensorboard_experiment_with_metadata` interceptor runs after the + `post_create_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_experiment` will be passed to + `post_create_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_run( self, request: tensorboard_service.CreateTensorboardRunRequest, @@ -486,12 +613,37 @@ def post_create_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for create_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_run` interceptor runs + before the `post_create_tensorboard_run_with_metadata` interceptor. """ return response + def post_create_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_run` interceptor. + When both interceptors are used, this `post_create_tensorboard_run_with_metadata` interceptor runs after the + `post_create_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_run` will be passed to + `post_create_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_time_series( self, request: tensorboard_service.CreateTensorboardTimeSeriesRequest, @@ -512,12 +664,38 @@ def post_create_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_time_series` interceptor runs + before the `post_create_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_create_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_time_series` will be passed to + `post_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard( self, request: tensorboard_service.DeleteTensorboardRequest, @@ -538,12 +716,35 @@ def post_delete_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard` interceptor runs + before the `post_delete_tensorboard_with_metadata` interceptor. """ return response + def post_delete_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard` interceptor. + When both interceptors are used, this `post_delete_tensorboard_with_metadata` interceptor runs after the + `post_delete_tensorboard` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard` will be passed to + `post_delete_tensorboard_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_experiment( self, request: tensorboard_service.DeleteTensorboardExperimentRequest, @@ -564,12 +765,35 @@ def post_delete_tensorboard_experiment( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_experiment` interceptor runs + before the `post_delete_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_delete_tensorboard_experiment_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_delete_tensorboard_experiment_with_metadata` interceptor runs after the + `post_delete_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_experiment` will be passed to + `post_delete_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_run( self, request: tensorboard_service.DeleteTensorboardRunRequest, @@ -590,12 +814,35 @@ def post_delete_tensorboard_run( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_run` interceptor runs + before the `post_delete_tensorboard_run_with_metadata` interceptor. """ return response + def post_delete_tensorboard_run_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_run` interceptor. + When both interceptors are used, this `post_delete_tensorboard_run_with_metadata` interceptor runs after the + `post_delete_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_run` will be passed to + `post_delete_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_time_series( self, request: tensorboard_service.DeleteTensorboardTimeSeriesRequest, @@ -616,12 +863,35 @@ def post_delete_tensorboard_time_series( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_time_series` interceptor runs + before the `post_delete_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_delete_tensorboard_time_series_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_delete_tensorboard_time_series_with_metadata` interceptor runs after the + `post_delete_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_time_series` will be passed to + `post_delete_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_export_tensorboard_time_series_data( self, request: tensorboard_service.ExportTensorboardTimeSeriesDataRequest, @@ -642,12 +912,38 @@ def post_export_tensorboard_time_series_data( ) -> tensorboard_service.ExportTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for export_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_export_tensorboard_time_series_data` interceptor runs + before the `post_export_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_export_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for export_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_export_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_export_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_export_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_export_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_export_tensorboard_time_series_data` will be passed to + `post_export_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard( self, request: tensorboard_service.GetTensorboardRequest, @@ -668,12 +964,35 @@ def post_get_tensorboard( ) -> tensorboard.Tensorboard: """Post-rpc interceptor for get_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard` interceptor runs + before the `post_get_tensorboard_with_metadata` interceptor. """ return response + def post_get_tensorboard_with_metadata( + self, + response: tensorboard.Tensorboard, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard.Tensorboard, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_with_metadata` + interceptor in new development instead of the `post_get_tensorboard` interceptor. + When both interceptors are used, this `post_get_tensorboard_with_metadata` interceptor runs after the + `post_get_tensorboard` interceptor. The (possibly modified) response returned by + `post_get_tensorboard` will be passed to + `post_get_tensorboard_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_experiment( self, request: tensorboard_service.GetTensorboardExperimentRequest, @@ -694,12 +1013,38 @@ def post_get_tensorboard_experiment( ) -> tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for get_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_experiment` interceptor runs + before the `post_get_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_get_tensorboard_experiment_with_metadata( + self, + response: tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_get_tensorboard_experiment_with_metadata` interceptor runs after the + `post_get_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_experiment` will be passed to + `post_get_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_run( self, request: tensorboard_service.GetTensorboardRunRequest, @@ -720,12 +1065,35 @@ def post_get_tensorboard_run( ) -> tensorboard_run.TensorboardRun: """Post-rpc interceptor for get_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_run` interceptor runs + before the `post_get_tensorboard_run_with_metadata` interceptor. """ return response + def post_get_tensorboard_run_with_metadata( + self, + response: tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_run` interceptor. + When both interceptors are used, this `post_get_tensorboard_run_with_metadata` interceptor runs after the + `post_get_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_run` will be passed to + `post_get_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_time_series( self, request: tensorboard_service.GetTensorboardTimeSeriesRequest, @@ -746,12 +1114,38 @@ def post_get_tensorboard_time_series( ) -> tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for get_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_time_series` interceptor runs + before the `post_get_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_get_tensorboard_time_series_with_metadata( + self, + response: tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_get_tensorboard_time_series_with_metadata` interceptor runs after the + `post_get_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_time_series` will be passed to + `post_get_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_experiments( self, request: tensorboard_service.ListTensorboardExperimentsRequest, @@ -772,12 +1166,38 @@ def post_list_tensorboard_experiments( ) -> tensorboard_service.ListTensorboardExperimentsResponse: """Post-rpc interceptor for list_tensorboard_experiments - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_experiments_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_experiments` interceptor runs + before the `post_list_tensorboard_experiments_with_metadata` interceptor. """ return response + def post_list_tensorboard_experiments_with_metadata( + self, + response: tensorboard_service.ListTensorboardExperimentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardExperimentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_experiments + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_experiments_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_experiments` interceptor. + When both interceptors are used, this `post_list_tensorboard_experiments_with_metadata` interceptor runs after the + `post_list_tensorboard_experiments` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_experiments` will be passed to + `post_list_tensorboard_experiments_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_runs( self, request: tensorboard_service.ListTensorboardRunsRequest, @@ -798,12 +1218,38 @@ def post_list_tensorboard_runs( ) -> tensorboard_service.ListTensorboardRunsResponse: """Post-rpc interceptor for list_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_runs` interceptor runs + before the `post_list_tensorboard_runs_with_metadata` interceptor. """ return response + def post_list_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.ListTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_runs` interceptor. + When both interceptors are used, this `post_list_tensorboard_runs_with_metadata` interceptor runs after the + `post_list_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_runs` will be passed to + `post_list_tensorboard_runs_with_metadata`. + """ + return response, metadata + def pre_list_tensorboards( self, request: tensorboard_service.ListTensorboardsRequest, @@ -824,12 +1270,38 @@ def post_list_tensorboards( ) -> tensorboard_service.ListTensorboardsResponse: """Post-rpc interceptor for list_tensorboards - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboards_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboards` interceptor runs + before the `post_list_tensorboards_with_metadata` interceptor. """ return response + def post_list_tensorboards_with_metadata( + self, + response: tensorboard_service.ListTensorboardsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboards + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboards_with_metadata` + interceptor in new development instead of the `post_list_tensorboards` interceptor. + When both interceptors are used, this `post_list_tensorboards_with_metadata` interceptor runs after the + `post_list_tensorboards` interceptor. The (possibly modified) response returned by + `post_list_tensorboards` will be passed to + `post_list_tensorboards_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_time_series( self, request: tensorboard_service.ListTensorboardTimeSeriesRequest, @@ -850,12 +1322,38 @@ def post_list_tensorboard_time_series( ) -> tensorboard_service.ListTensorboardTimeSeriesResponse: """Post-rpc interceptor for list_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_time_series` interceptor runs + before the `post_list_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_list_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.ListTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_list_tensorboard_time_series_with_metadata` interceptor runs after the + `post_list_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_time_series` will be passed to + `post_list_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_blob_data( self, request: tensorboard_service.ReadTensorboardBlobDataRequest, @@ -876,12 +1374,37 @@ def post_read_tensorboard_blob_data( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for read_tensorboard_blob_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_blob_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_blob_data` interceptor runs + before the `post_read_tensorboard_blob_data_with_metadata` interceptor. """ return response + def post_read_tensorboard_blob_data_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for read_tensorboard_blob_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_blob_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_blob_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_blob_data_with_metadata` interceptor runs after the + `post_read_tensorboard_blob_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_blob_data` will be passed to + `post_read_tensorboard_blob_data_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_size( self, request: tensorboard_service.ReadTensorboardSizeRequest, @@ -902,12 +1425,38 @@ def post_read_tensorboard_size( ) -> tensorboard_service.ReadTensorboardSizeResponse: """Post-rpc interceptor for read_tensorboard_size - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_size_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_size` interceptor runs + before the `post_read_tensorboard_size_with_metadata` interceptor. """ return response + def post_read_tensorboard_size_with_metadata( + self, + response: tensorboard_service.ReadTensorboardSizeResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardSizeResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_size + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_size_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_size` interceptor. + When both interceptors are used, this `post_read_tensorboard_size_with_metadata` interceptor runs after the + `post_read_tensorboard_size` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_size` will be passed to + `post_read_tensorboard_size_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_time_series_data( self, request: tensorboard_service.ReadTensorboardTimeSeriesDataRequest, @@ -928,12 +1477,38 @@ def post_read_tensorboard_time_series_data( ) -> tensorboard_service.ReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_time_series_data` interceptor runs + before the `post_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_time_series_data` will be passed to + `post_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_usage( self, request: tensorboard_service.ReadTensorboardUsageRequest, @@ -954,12 +1529,38 @@ def post_read_tensorboard_usage( ) -> tensorboard_service.ReadTensorboardUsageResponse: """Post-rpc interceptor for read_tensorboard_usage - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_usage_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_usage` interceptor runs + before the `post_read_tensorboard_usage_with_metadata` interceptor. """ return response + def post_read_tensorboard_usage_with_metadata( + self, + response: tensorboard_service.ReadTensorboardUsageResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardUsageResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_usage + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_usage_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_usage` interceptor. + When both interceptors are used, this `post_read_tensorboard_usage_with_metadata` interceptor runs after the + `post_read_tensorboard_usage` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_usage` will be passed to + `post_read_tensorboard_usage_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard( self, request: tensorboard_service.UpdateTensorboardRequest, @@ -980,12 +1581,35 @@ def post_update_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard` interceptor runs + before the `post_update_tensorboard_with_metadata` interceptor. """ return response + def post_update_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_with_metadata` + interceptor in new development instead of the `post_update_tensorboard` interceptor. + When both interceptors are used, this `post_update_tensorboard_with_metadata` interceptor runs after the + `post_update_tensorboard` interceptor. The (possibly modified) response returned by + `post_update_tensorboard` will be passed to + `post_update_tensorboard_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_experiment( self, request: tensorboard_service.UpdateTensorboardExperimentRequest, @@ -1006,12 +1630,38 @@ def post_update_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for update_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_experiment` interceptor runs + before the `post_update_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_update_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_update_tensorboard_experiment_with_metadata` interceptor runs after the + `post_update_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_experiment` will be passed to + `post_update_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_run( self, request: tensorboard_service.UpdateTensorboardRunRequest, @@ -1032,12 +1682,37 @@ def post_update_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for update_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_run` interceptor runs + before the `post_update_tensorboard_run_with_metadata` interceptor. """ return response + def post_update_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_run` interceptor. + When both interceptors are used, this `post_update_tensorboard_run_with_metadata` interceptor runs after the + `post_update_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_run` will be passed to + `post_update_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_time_series( self, request: tensorboard_service.UpdateTensorboardTimeSeriesRequest, @@ -1058,12 +1733,38 @@ def post_update_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for update_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_time_series` interceptor runs + before the `post_update_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_update_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_update_tensorboard_time_series_with_metadata` interceptor runs after the + `post_update_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_time_series` will be passed to + `post_update_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_write_tensorboard_experiment_data( self, request: tensorboard_service.WriteTensorboardExperimentDataRequest, @@ -1084,12 +1785,38 @@ def post_write_tensorboard_experiment_data( ) -> tensorboard_service.WriteTensorboardExperimentDataResponse: """Post-rpc interceptor for write_tensorboard_experiment_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_experiment_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_experiment_data` interceptor runs + before the `post_write_tensorboard_experiment_data_with_metadata` interceptor. """ return response + def post_write_tensorboard_experiment_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardExperimentDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardExperimentDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_experiment_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_experiment_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_experiment_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_experiment_data_with_metadata` interceptor runs after the + `post_write_tensorboard_experiment_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_experiment_data` will be passed to + `post_write_tensorboard_experiment_data_with_metadata`. + """ + return response, metadata + def pre_write_tensorboard_run_data( self, request: tensorboard_service.WriteTensorboardRunDataRequest, @@ -1110,12 +1837,38 @@ def post_write_tensorboard_run_data( ) -> tensorboard_service.WriteTensorboardRunDataResponse: """Post-rpc interceptor for write_tensorboard_run_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_run_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_run_data` interceptor runs + before the `post_write_tensorboard_run_data_with_metadata` interceptor. """ return response + def post_write_tensorboard_run_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardRunDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardRunDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_run_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_run_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_run_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_run_data_with_metadata` interceptor runs after the + `post_write_tensorboard_run_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_run_data` will be passed to + `post_write_tensorboard_run_data_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3492,6 +4245,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_create_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3657,6 +4417,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3814,6 +4581,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3968,6 +4742,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4124,6 +4902,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4285,6 +5070,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4441,6 +5230,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4591,6 +5387,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4738,6 +5538,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4886,6 +5693,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5033,6 +5844,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5194,6 +6012,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_export_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5345,6 +6170,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5497,6 +6326,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5648,6 +6481,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5798,6 +6635,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5948,6 +6789,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_experiments(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_experiments_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6101,6 +6946,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6252,6 +7101,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboards(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboards_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6402,6 +7255,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6556,6 +7413,10 @@ def __call__( ) resp = self._interceptor.post_read_tensorboard_blob_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_blob_data_with_metadata( + resp, response_metadata + ) return resp class _ReadTensorboardSize( @@ -6684,6 +7545,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_size(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_size_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6836,6 +7701,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6987,6 +7859,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_usage(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_usage_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7143,6 +8019,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7299,6 +8179,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7460,6 +8347,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7616,6 +8507,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7778,6 +8676,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_tensorboard_experiment_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_write_tensorboard_experiment_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7936,6 +8841,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_tensorboard_run_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_write_tensorboard_run_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest_asyncio.py index 460f8288c3..a3dcce5a89 100644 --- a/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/tensorboard_service/transports/rest_asyncio.py @@ -374,12 +374,38 @@ async def post_batch_create_tensorboard_runs( ) -> tensorboard_service.BatchCreateTensorboardRunsResponse: """Post-rpc interceptor for batch_create_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_runs` interceptor runs + before the `post_batch_create_tensorboard_runs_with_metadata` interceptor. """ return response + async def post_batch_create_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_runs` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_runs_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_runs` will be passed to + `post_batch_create_tensorboard_runs_with_metadata`. + """ + return response, metadata + async def pre_batch_create_tensorboard_time_series( self, request: tensorboard_service.BatchCreateTensorboardTimeSeriesRequest, @@ -400,12 +426,38 @@ async def post_batch_create_tensorboard_time_series( ) -> tensorboard_service.BatchCreateTensorboardTimeSeriesResponse: """Post-rpc interceptor for batch_create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_time_series` interceptor runs + before the `post_batch_create_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_batch_create_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_time_series` will be passed to + `post_batch_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_batch_read_tensorboard_time_series_data( self, request: tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest, @@ -426,12 +478,38 @@ async def post_batch_read_tensorboard_time_series_data( ) -> tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for batch_read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_tensorboard_time_series_data` interceptor runs + before the `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_batch_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_batch_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_batch_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_batch_read_tensorboard_time_series_data` will be passed to + `post_batch_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard( self, request: tensorboard_service.CreateTensorboardRequest, @@ -452,12 +530,35 @@ async def post_create_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard` interceptor runs + before the `post_create_tensorboard_with_metadata` interceptor. """ return response + async def post_create_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_with_metadata` + interceptor in new development instead of the `post_create_tensorboard` interceptor. + When both interceptors are used, this `post_create_tensorboard_with_metadata` interceptor runs after the + `post_create_tensorboard` interceptor. The (possibly modified) response returned by + `post_create_tensorboard` will be passed to + `post_create_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_experiment( self, request: tensorboard_service.CreateTensorboardExperimentRequest, @@ -478,12 +579,38 @@ async def post_create_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for create_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_experiment` interceptor runs + before the `post_create_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_create_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_create_tensorboard_experiment_with_metadata` interceptor runs after the + `post_create_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_experiment` will be passed to + `post_create_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_run( self, request: tensorboard_service.CreateTensorboardRunRequest, @@ -504,12 +631,37 @@ async def post_create_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for create_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_run` interceptor runs + before the `post_create_tensorboard_run_with_metadata` interceptor. """ return response + async def post_create_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_run` interceptor. + When both interceptors are used, this `post_create_tensorboard_run_with_metadata` interceptor runs after the + `post_create_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_run` will be passed to + `post_create_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_time_series( self, request: tensorboard_service.CreateTensorboardTimeSeriesRequest, @@ -530,12 +682,38 @@ async def post_create_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_time_series` interceptor runs + before the `post_create_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_create_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_time_series` will be passed to + `post_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard( self, request: tensorboard_service.DeleteTensorboardRequest, @@ -556,12 +734,35 @@ async def post_delete_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard` interceptor runs + before the `post_delete_tensorboard_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard` interceptor. + When both interceptors are used, this `post_delete_tensorboard_with_metadata` interceptor runs after the + `post_delete_tensorboard` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard` will be passed to + `post_delete_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_experiment( self, request: tensorboard_service.DeleteTensorboardExperimentRequest, @@ -582,12 +783,35 @@ async def post_delete_tensorboard_experiment( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_experiment` interceptor runs + before the `post_delete_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_experiment_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_delete_tensorboard_experiment_with_metadata` interceptor runs after the + `post_delete_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_experiment` will be passed to + `post_delete_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_run( self, request: tensorboard_service.DeleteTensorboardRunRequest, @@ -608,12 +832,35 @@ async def post_delete_tensorboard_run( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_run` interceptor runs + before the `post_delete_tensorboard_run_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_run_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_run` interceptor. + When both interceptors are used, this `post_delete_tensorboard_run_with_metadata` interceptor runs after the + `post_delete_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_run` will be passed to + `post_delete_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_time_series( self, request: tensorboard_service.DeleteTensorboardTimeSeriesRequest, @@ -634,12 +881,35 @@ async def post_delete_tensorboard_time_series( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_time_series` interceptor runs + before the `post_delete_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_time_series_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_delete_tensorboard_time_series_with_metadata` interceptor runs after the + `post_delete_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_time_series` will be passed to + `post_delete_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_export_tensorboard_time_series_data( self, request: tensorboard_service.ExportTensorboardTimeSeriesDataRequest, @@ -660,12 +930,38 @@ async def post_export_tensorboard_time_series_data( ) -> tensorboard_service.ExportTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for export_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_export_tensorboard_time_series_data` interceptor runs + before the `post_export_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_export_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for export_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_export_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_export_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_export_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_export_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_export_tensorboard_time_series_data` will be passed to + `post_export_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard( self, request: tensorboard_service.GetTensorboardRequest, @@ -686,12 +982,35 @@ async def post_get_tensorboard( ) -> tensorboard.Tensorboard: """Post-rpc interceptor for get_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard` interceptor runs + before the `post_get_tensorboard_with_metadata` interceptor. """ return response + async def post_get_tensorboard_with_metadata( + self, + response: tensorboard.Tensorboard, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard.Tensorboard, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_with_metadata` + interceptor in new development instead of the `post_get_tensorboard` interceptor. + When both interceptors are used, this `post_get_tensorboard_with_metadata` interceptor runs after the + `post_get_tensorboard` interceptor. The (possibly modified) response returned by + `post_get_tensorboard` will be passed to + `post_get_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_experiment( self, request: tensorboard_service.GetTensorboardExperimentRequest, @@ -712,12 +1031,38 @@ async def post_get_tensorboard_experiment( ) -> tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for get_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_experiment` interceptor runs + before the `post_get_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_get_tensorboard_experiment_with_metadata( + self, + response: tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_get_tensorboard_experiment_with_metadata` interceptor runs after the + `post_get_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_experiment` will be passed to + `post_get_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_run( self, request: tensorboard_service.GetTensorboardRunRequest, @@ -738,12 +1083,35 @@ async def post_get_tensorboard_run( ) -> tensorboard_run.TensorboardRun: """Post-rpc interceptor for get_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_run` interceptor runs + before the `post_get_tensorboard_run_with_metadata` interceptor. """ return response + async def post_get_tensorboard_run_with_metadata( + self, + response: tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_run` interceptor. + When both interceptors are used, this `post_get_tensorboard_run_with_metadata` interceptor runs after the + `post_get_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_run` will be passed to + `post_get_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_time_series( self, request: tensorboard_service.GetTensorboardTimeSeriesRequest, @@ -764,12 +1132,38 @@ async def post_get_tensorboard_time_series( ) -> tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for get_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_time_series` interceptor runs + before the `post_get_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_get_tensorboard_time_series_with_metadata( + self, + response: tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_get_tensorboard_time_series_with_metadata` interceptor runs after the + `post_get_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_time_series` will be passed to + `post_get_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_experiments( self, request: tensorboard_service.ListTensorboardExperimentsRequest, @@ -790,12 +1184,38 @@ async def post_list_tensorboard_experiments( ) -> tensorboard_service.ListTensorboardExperimentsResponse: """Post-rpc interceptor for list_tensorboard_experiments - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_experiments_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_experiments` interceptor runs + before the `post_list_tensorboard_experiments_with_metadata` interceptor. """ return response + async def post_list_tensorboard_experiments_with_metadata( + self, + response: tensorboard_service.ListTensorboardExperimentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardExperimentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_experiments + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_experiments_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_experiments` interceptor. + When both interceptors are used, this `post_list_tensorboard_experiments_with_metadata` interceptor runs after the + `post_list_tensorboard_experiments` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_experiments` will be passed to + `post_list_tensorboard_experiments_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_runs( self, request: tensorboard_service.ListTensorboardRunsRequest, @@ -816,12 +1236,38 @@ async def post_list_tensorboard_runs( ) -> tensorboard_service.ListTensorboardRunsResponse: """Post-rpc interceptor for list_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_runs` interceptor runs + before the `post_list_tensorboard_runs_with_metadata` interceptor. """ return response + async def post_list_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.ListTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_runs` interceptor. + When both interceptors are used, this `post_list_tensorboard_runs_with_metadata` interceptor runs after the + `post_list_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_runs` will be passed to + `post_list_tensorboard_runs_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboards( self, request: tensorboard_service.ListTensorboardsRequest, @@ -842,12 +1288,38 @@ async def post_list_tensorboards( ) -> tensorboard_service.ListTensorboardsResponse: """Post-rpc interceptor for list_tensorboards - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboards_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboards` interceptor runs + before the `post_list_tensorboards_with_metadata` interceptor. """ return response + async def post_list_tensorboards_with_metadata( + self, + response: tensorboard_service.ListTensorboardsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboards + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboards_with_metadata` + interceptor in new development instead of the `post_list_tensorboards` interceptor. + When both interceptors are used, this `post_list_tensorboards_with_metadata` interceptor runs after the + `post_list_tensorboards` interceptor. The (possibly modified) response returned by + `post_list_tensorboards` will be passed to + `post_list_tensorboards_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_time_series( self, request: tensorboard_service.ListTensorboardTimeSeriesRequest, @@ -868,12 +1340,38 @@ async def post_list_tensorboard_time_series( ) -> tensorboard_service.ListTensorboardTimeSeriesResponse: """Post-rpc interceptor for list_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_time_series` interceptor runs + before the `post_list_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_list_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.ListTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_list_tensorboard_time_series_with_metadata` interceptor runs after the + `post_list_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_time_series` will be passed to + `post_list_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_blob_data( self, request: tensorboard_service.ReadTensorboardBlobDataRequest, @@ -894,12 +1392,38 @@ async def post_read_tensorboard_blob_data( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for read_tensorboard_blob_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_blob_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_blob_data` interceptor runs + before the `post_read_tensorboard_blob_data_with_metadata` interceptor. """ return response + async def post_read_tensorboard_blob_data_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_blob_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_blob_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_blob_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_blob_data_with_metadata` interceptor runs after the + `post_read_tensorboard_blob_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_blob_data` will be passed to + `post_read_tensorboard_blob_data_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_size( self, request: tensorboard_service.ReadTensorboardSizeRequest, @@ -920,12 +1444,38 @@ async def post_read_tensorboard_size( ) -> tensorboard_service.ReadTensorboardSizeResponse: """Post-rpc interceptor for read_tensorboard_size - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_size_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_size` interceptor runs + before the `post_read_tensorboard_size_with_metadata` interceptor. """ return response + async def post_read_tensorboard_size_with_metadata( + self, + response: tensorboard_service.ReadTensorboardSizeResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardSizeResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_size + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_size_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_size` interceptor. + When both interceptors are used, this `post_read_tensorboard_size_with_metadata` interceptor runs after the + `post_read_tensorboard_size` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_size` will be passed to + `post_read_tensorboard_size_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_time_series_data( self, request: tensorboard_service.ReadTensorboardTimeSeriesDataRequest, @@ -946,12 +1496,38 @@ async def post_read_tensorboard_time_series_data( ) -> tensorboard_service.ReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_time_series_data` interceptor runs + before the `post_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_time_series_data` will be passed to + `post_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_usage( self, request: tensorboard_service.ReadTensorboardUsageRequest, @@ -972,12 +1548,38 @@ async def post_read_tensorboard_usage( ) -> tensorboard_service.ReadTensorboardUsageResponse: """Post-rpc interceptor for read_tensorboard_usage - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_usage_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_usage` interceptor runs + before the `post_read_tensorboard_usage_with_metadata` interceptor. """ return response + async def post_read_tensorboard_usage_with_metadata( + self, + response: tensorboard_service.ReadTensorboardUsageResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardUsageResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_usage + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_usage_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_usage` interceptor. + When both interceptors are used, this `post_read_tensorboard_usage_with_metadata` interceptor runs after the + `post_read_tensorboard_usage` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_usage` will be passed to + `post_read_tensorboard_usage_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard( self, request: tensorboard_service.UpdateTensorboardRequest, @@ -998,12 +1600,35 @@ async def post_update_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard` interceptor runs + before the `post_update_tensorboard_with_metadata` interceptor. """ return response + async def post_update_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_with_metadata` + interceptor in new development instead of the `post_update_tensorboard` interceptor. + When both interceptors are used, this `post_update_tensorboard_with_metadata` interceptor runs after the + `post_update_tensorboard` interceptor. The (possibly modified) response returned by + `post_update_tensorboard` will be passed to + `post_update_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_experiment( self, request: tensorboard_service.UpdateTensorboardExperimentRequest, @@ -1024,12 +1649,38 @@ async def post_update_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for update_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_experiment` interceptor runs + before the `post_update_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_update_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_update_tensorboard_experiment_with_metadata` interceptor runs after the + `post_update_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_experiment` will be passed to + `post_update_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_run( self, request: tensorboard_service.UpdateTensorboardRunRequest, @@ -1050,12 +1701,37 @@ async def post_update_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for update_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_run` interceptor runs + before the `post_update_tensorboard_run_with_metadata` interceptor. """ return response + async def post_update_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_run` interceptor. + When both interceptors are used, this `post_update_tensorboard_run_with_metadata` interceptor runs after the + `post_update_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_run` will be passed to + `post_update_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_time_series( self, request: tensorboard_service.UpdateTensorboardTimeSeriesRequest, @@ -1076,12 +1752,38 @@ async def post_update_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for update_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_time_series` interceptor runs + before the `post_update_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_update_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_update_tensorboard_time_series_with_metadata` interceptor runs after the + `post_update_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_time_series` will be passed to + `post_update_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_write_tensorboard_experiment_data( self, request: tensorboard_service.WriteTensorboardExperimentDataRequest, @@ -1102,12 +1804,38 @@ async def post_write_tensorboard_experiment_data( ) -> tensorboard_service.WriteTensorboardExperimentDataResponse: """Post-rpc interceptor for write_tensorboard_experiment_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_experiment_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_experiment_data` interceptor runs + before the `post_write_tensorboard_experiment_data_with_metadata` interceptor. """ return response + async def post_write_tensorboard_experiment_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardExperimentDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardExperimentDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_experiment_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_experiment_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_experiment_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_experiment_data_with_metadata` interceptor runs after the + `post_write_tensorboard_experiment_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_experiment_data` will be passed to + `post_write_tensorboard_experiment_data_with_metadata`. + """ + return response, metadata + async def pre_write_tensorboard_run_data( self, request: tensorboard_service.WriteTensorboardRunDataRequest, @@ -1128,12 +1856,38 @@ async def post_write_tensorboard_run_data( ) -> tensorboard_service.WriteTensorboardRunDataResponse: """Post-rpc interceptor for write_tensorboard_run_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_run_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_run_data` interceptor runs + before the `post_write_tensorboard_run_data_with_metadata` interceptor. """ return response + async def post_write_tensorboard_run_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardRunDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardRunDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_run_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_run_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_run_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_run_data_with_metadata` interceptor runs after the + `post_write_tensorboard_run_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_run_data` will be passed to + `post_write_tensorboard_run_data_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1797,6 +2551,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_create_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1970,6 +2731,13 @@ async def __call__( resp = await self._interceptor.post_batch_create_tensorboard_time_series( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2135,6 +2903,13 @@ async def __call__( resp = await self._interceptor.post_batch_read_tensorboard_time_series_data( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2297,6 +3072,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2464,6 +3243,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2629,6 +3415,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2796,6 +3586,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2954,6 +3751,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3114,6 +3915,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3268,6 +4076,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3428,6 +4240,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3597,6 +4416,13 @@ async def __call__( resp = await self._interceptor.post_export_tensorboard_time_series_data( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_export_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3756,6 +4582,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3912,6 +4742,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4069,6 +4906,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4223,6 +5064,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4384,6 +5232,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_experiments(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_tensorboard_experiments_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4541,6 +5396,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4698,6 +5557,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboards(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tensorboards_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4859,6 +5722,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5016,6 +5886,13 @@ async def __call__( response, tensorboard_service.ReadTensorboardBlobDataResponse ) resp = await self._interceptor.post_read_tensorboard_blob_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_read_tensorboard_blob_data_with_metadata( + resp, response_metadata + ) return resp class _ReadTensorboardSize( @@ -5147,6 +6024,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_size(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_tensorboard_size_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5310,6 +6191,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5465,6 +6353,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_usage(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_tensorboard_usage_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5629,6 +6521,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5796,6 +6692,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5961,6 +6864,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6128,6 +7035,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6299,6 +7213,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_tensorboard_experiment_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_write_tensorboard_experiment_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6461,6 +7382,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_tensorboard_run_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_write_tensorboard_run_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/client.py b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/client.py index 184f1a230b..e4007a9a8d 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/client.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -599,6 +601,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2134,16 +2163,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2189,16 +2222,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2361,16 +2398,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2482,16 +2523,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2604,16 +2649,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2664,16 +2713,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2719,16 +2772,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2774,16 +2831,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest.py b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest.py index 73e4cedfec..3689a02dbd 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest.py @@ -186,12 +186,35 @@ def post_create_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_create_rag_corpus` interceptor runs + before the `post_create_rag_corpus_with_metadata` interceptor. """ return response + def post_create_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_create_rag_corpus_with_metadata` + interceptor in new development instead of the `post_create_rag_corpus` interceptor. + When both interceptors are used, this `post_create_rag_corpus_with_metadata` interceptor runs after the + `post_create_rag_corpus` interceptor. The (possibly modified) response returned by + `post_create_rag_corpus` will be passed to + `post_create_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_delete_rag_corpus( self, request: vertex_rag_data_service.DeleteRagCorpusRequest, @@ -212,12 +235,35 @@ def post_delete_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_corpus` interceptor runs + before the `post_delete_rag_corpus_with_metadata` interceptor. """ return response + def post_delete_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_corpus_with_metadata` + interceptor in new development instead of the `post_delete_rag_corpus` interceptor. + When both interceptors are used, this `post_delete_rag_corpus_with_metadata` interceptor runs after the + `post_delete_rag_corpus` interceptor. The (possibly modified) response returned by + `post_delete_rag_corpus` will be passed to + `post_delete_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_delete_rag_file( self, request: vertex_rag_data_service.DeleteRagFileRequest, @@ -238,12 +284,35 @@ def post_delete_rag_file( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_file` interceptor runs + before the `post_delete_rag_file_with_metadata` interceptor. """ return response + def post_delete_rag_file_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_file_with_metadata` + interceptor in new development instead of the `post_delete_rag_file` interceptor. + When both interceptors are used, this `post_delete_rag_file_with_metadata` interceptor runs after the + `post_delete_rag_file` interceptor. The (possibly modified) response returned by + `post_delete_rag_file` will be passed to + `post_delete_rag_file_with_metadata`. + """ + return response, metadata + def pre_get_rag_corpus( self, request: vertex_rag_data_service.GetRagCorpusRequest, @@ -264,12 +333,35 @@ def post_get_rag_corpus( ) -> vertex_rag_data.RagCorpus: """Post-rpc interceptor for get_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_corpus` interceptor runs + before the `post_get_rag_corpus_with_metadata` interceptor. """ return response + def post_get_rag_corpus_with_metadata( + self, + response: vertex_rag_data.RagCorpus, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagCorpus, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_corpus_with_metadata` + interceptor in new development instead of the `post_get_rag_corpus` interceptor. + When both interceptors are used, this `post_get_rag_corpus_with_metadata` interceptor runs after the + `post_get_rag_corpus` interceptor. The (possibly modified) response returned by + `post_get_rag_corpus` will be passed to + `post_get_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_get_rag_file( self, request: vertex_rag_data_service.GetRagFileRequest, @@ -290,12 +382,35 @@ def post_get_rag_file( ) -> vertex_rag_data.RagFile: """Post-rpc interceptor for get_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_file` interceptor runs + before the `post_get_rag_file_with_metadata` interceptor. """ return response + def post_get_rag_file_with_metadata( + self, + response: vertex_rag_data.RagFile, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagFile, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_file_with_metadata` + interceptor in new development instead of the `post_get_rag_file` interceptor. + When both interceptors are used, this `post_get_rag_file_with_metadata` interceptor runs after the + `post_get_rag_file` interceptor. The (possibly modified) response returned by + `post_get_rag_file` will be passed to + `post_get_rag_file_with_metadata`. + """ + return response, metadata + def pre_import_rag_files( self, request: vertex_rag_data_service.ImportRagFilesRequest, @@ -316,12 +431,35 @@ def post_import_rag_files( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_import_rag_files` interceptor runs + before the `post_import_rag_files_with_metadata` interceptor. """ return response + def post_import_rag_files_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_import_rag_files_with_metadata` + interceptor in new development instead of the `post_import_rag_files` interceptor. + When both interceptors are used, this `post_import_rag_files_with_metadata` interceptor runs after the + `post_import_rag_files` interceptor. The (possibly modified) response returned by + `post_import_rag_files` will be passed to + `post_import_rag_files_with_metadata`. + """ + return response, metadata + def pre_list_rag_corpora( self, request: vertex_rag_data_service.ListRagCorporaRequest, @@ -342,12 +480,38 @@ def post_list_rag_corpora( ) -> vertex_rag_data_service.ListRagCorporaResponse: """Post-rpc interceptor for list_rag_corpora - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_corpora_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_corpora` interceptor runs + before the `post_list_rag_corpora_with_metadata` interceptor. """ return response + def post_list_rag_corpora_with_metadata( + self, + response: vertex_rag_data_service.ListRagCorporaResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagCorporaResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_corpora + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_corpora_with_metadata` + interceptor in new development instead of the `post_list_rag_corpora` interceptor. + When both interceptors are used, this `post_list_rag_corpora_with_metadata` interceptor runs after the + `post_list_rag_corpora` interceptor. The (possibly modified) response returned by + `post_list_rag_corpora` will be passed to + `post_list_rag_corpora_with_metadata`. + """ + return response, metadata + def pre_list_rag_files( self, request: vertex_rag_data_service.ListRagFilesRequest, @@ -368,12 +532,38 @@ def post_list_rag_files( ) -> vertex_rag_data_service.ListRagFilesResponse: """Post-rpc interceptor for list_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_files` interceptor runs + before the `post_list_rag_files_with_metadata` interceptor. """ return response + def post_list_rag_files_with_metadata( + self, + response: vertex_rag_data_service.ListRagFilesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagFilesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_files_with_metadata` + interceptor in new development instead of the `post_list_rag_files` interceptor. + When both interceptors are used, this `post_list_rag_files_with_metadata` interceptor runs after the + `post_list_rag_files` interceptor. The (possibly modified) response returned by + `post_list_rag_files` will be passed to + `post_list_rag_files_with_metadata`. + """ + return response, metadata + def pre_update_rag_corpus( self, request: vertex_rag_data_service.UpdateRagCorpusRequest, @@ -394,12 +584,35 @@ def post_update_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_update_rag_corpus` interceptor runs + before the `post_update_rag_corpus_with_metadata` interceptor. """ return response + def post_update_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_update_rag_corpus_with_metadata` + interceptor in new development instead of the `post_update_rag_corpus` interceptor. + When both interceptors are used, this `post_update_rag_corpus_with_metadata` interceptor runs after the + `post_update_rag_corpus` interceptor. The (possibly modified) response returned by + `post_update_rag_corpus` will be passed to + `post_update_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_upload_rag_file( self, request: vertex_rag_data_service.UploadRagFileRequest, @@ -420,12 +633,38 @@ def post_upload_rag_file( ) -> vertex_rag_data_service.UploadRagFileResponse: """Post-rpc interceptor for upload_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_upload_rag_file` interceptor runs + before the `post_upload_rag_file_with_metadata` interceptor. """ return response + def post_upload_rag_file_with_metadata( + self, + response: vertex_rag_data_service.UploadRagFileResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.UploadRagFileResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for upload_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_upload_rag_file_with_metadata` + interceptor in new development instead of the `post_upload_rag_file` interceptor. + When both interceptors are used, this `post_upload_rag_file_with_metadata` interceptor runs after the + `post_upload_rag_file` interceptor. The (possibly modified) response returned by + `post_upload_rag_file` will be passed to + `post_upload_rag_file_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2800,6 +3039,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2946,6 +3189,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3090,6 +3337,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3236,6 +3487,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3381,6 +3636,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3533,6 +3792,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3680,6 +3943,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_rag_corpora(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_rag_corpora_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3827,6 +4094,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3981,6 +4252,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4132,6 +4407,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_upload_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upload_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest_asyncio.py index 499e218620..3a50b226cb 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_data_service/transports/rest_asyncio.py @@ -204,12 +204,35 @@ async def post_create_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_create_rag_corpus` interceptor runs + before the `post_create_rag_corpus_with_metadata` interceptor. """ return response + async def post_create_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_create_rag_corpus_with_metadata` + interceptor in new development instead of the `post_create_rag_corpus` interceptor. + When both interceptors are used, this `post_create_rag_corpus_with_metadata` interceptor runs after the + `post_create_rag_corpus` interceptor. The (possibly modified) response returned by + `post_create_rag_corpus` will be passed to + `post_create_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_delete_rag_corpus( self, request: vertex_rag_data_service.DeleteRagCorpusRequest, @@ -230,12 +253,35 @@ async def post_delete_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_corpus` interceptor runs + before the `post_delete_rag_corpus_with_metadata` interceptor. """ return response + async def post_delete_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_corpus_with_metadata` + interceptor in new development instead of the `post_delete_rag_corpus` interceptor. + When both interceptors are used, this `post_delete_rag_corpus_with_metadata` interceptor runs after the + `post_delete_rag_corpus` interceptor. The (possibly modified) response returned by + `post_delete_rag_corpus` will be passed to + `post_delete_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_delete_rag_file( self, request: vertex_rag_data_service.DeleteRagFileRequest, @@ -256,12 +302,35 @@ async def post_delete_rag_file( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_file` interceptor runs + before the `post_delete_rag_file_with_metadata` interceptor. """ return response + async def post_delete_rag_file_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_file_with_metadata` + interceptor in new development instead of the `post_delete_rag_file` interceptor. + When both interceptors are used, this `post_delete_rag_file_with_metadata` interceptor runs after the + `post_delete_rag_file` interceptor. The (possibly modified) response returned by + `post_delete_rag_file` will be passed to + `post_delete_rag_file_with_metadata`. + """ + return response, metadata + async def pre_get_rag_corpus( self, request: vertex_rag_data_service.GetRagCorpusRequest, @@ -282,12 +351,35 @@ async def post_get_rag_corpus( ) -> vertex_rag_data.RagCorpus: """Post-rpc interceptor for get_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_corpus` interceptor runs + before the `post_get_rag_corpus_with_metadata` interceptor. """ return response + async def post_get_rag_corpus_with_metadata( + self, + response: vertex_rag_data.RagCorpus, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagCorpus, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_corpus_with_metadata` + interceptor in new development instead of the `post_get_rag_corpus` interceptor. + When both interceptors are used, this `post_get_rag_corpus_with_metadata` interceptor runs after the + `post_get_rag_corpus` interceptor. The (possibly modified) response returned by + `post_get_rag_corpus` will be passed to + `post_get_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_get_rag_file( self, request: vertex_rag_data_service.GetRagFileRequest, @@ -308,12 +400,35 @@ async def post_get_rag_file( ) -> vertex_rag_data.RagFile: """Post-rpc interceptor for get_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_file` interceptor runs + before the `post_get_rag_file_with_metadata` interceptor. """ return response + async def post_get_rag_file_with_metadata( + self, + response: vertex_rag_data.RagFile, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagFile, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_file_with_metadata` + interceptor in new development instead of the `post_get_rag_file` interceptor. + When both interceptors are used, this `post_get_rag_file_with_metadata` interceptor runs after the + `post_get_rag_file` interceptor. The (possibly modified) response returned by + `post_get_rag_file` will be passed to + `post_get_rag_file_with_metadata`. + """ + return response, metadata + async def pre_import_rag_files( self, request: vertex_rag_data_service.ImportRagFilesRequest, @@ -334,12 +449,35 @@ async def post_import_rag_files( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_import_rag_files` interceptor runs + before the `post_import_rag_files_with_metadata` interceptor. """ return response + async def post_import_rag_files_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_import_rag_files_with_metadata` + interceptor in new development instead of the `post_import_rag_files` interceptor. + When both interceptors are used, this `post_import_rag_files_with_metadata` interceptor runs after the + `post_import_rag_files` interceptor. The (possibly modified) response returned by + `post_import_rag_files` will be passed to + `post_import_rag_files_with_metadata`. + """ + return response, metadata + async def pre_list_rag_corpora( self, request: vertex_rag_data_service.ListRagCorporaRequest, @@ -360,12 +498,38 @@ async def post_list_rag_corpora( ) -> vertex_rag_data_service.ListRagCorporaResponse: """Post-rpc interceptor for list_rag_corpora - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_corpora_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_corpora` interceptor runs + before the `post_list_rag_corpora_with_metadata` interceptor. """ return response + async def post_list_rag_corpora_with_metadata( + self, + response: vertex_rag_data_service.ListRagCorporaResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagCorporaResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_corpora + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_corpora_with_metadata` + interceptor in new development instead of the `post_list_rag_corpora` interceptor. + When both interceptors are used, this `post_list_rag_corpora_with_metadata` interceptor runs after the + `post_list_rag_corpora` interceptor. The (possibly modified) response returned by + `post_list_rag_corpora` will be passed to + `post_list_rag_corpora_with_metadata`. + """ + return response, metadata + async def pre_list_rag_files( self, request: vertex_rag_data_service.ListRagFilesRequest, @@ -386,12 +550,38 @@ async def post_list_rag_files( ) -> vertex_rag_data_service.ListRagFilesResponse: """Post-rpc interceptor for list_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_files` interceptor runs + before the `post_list_rag_files_with_metadata` interceptor. """ return response + async def post_list_rag_files_with_metadata( + self, + response: vertex_rag_data_service.ListRagFilesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagFilesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_files_with_metadata` + interceptor in new development instead of the `post_list_rag_files` interceptor. + When both interceptors are used, this `post_list_rag_files_with_metadata` interceptor runs after the + `post_list_rag_files` interceptor. The (possibly modified) response returned by + `post_list_rag_files` will be passed to + `post_list_rag_files_with_metadata`. + """ + return response, metadata + async def pre_update_rag_corpus( self, request: vertex_rag_data_service.UpdateRagCorpusRequest, @@ -412,12 +602,35 @@ async def post_update_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_update_rag_corpus` interceptor runs + before the `post_update_rag_corpus_with_metadata` interceptor. """ return response + async def post_update_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_update_rag_corpus_with_metadata` + interceptor in new development instead of the `post_update_rag_corpus` interceptor. + When both interceptors are used, this `post_update_rag_corpus_with_metadata` interceptor runs after the + `post_update_rag_corpus` interceptor. The (possibly modified) response returned by + `post_update_rag_corpus` will be passed to + `post_update_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_upload_rag_file( self, request: vertex_rag_data_service.UploadRagFileRequest, @@ -438,12 +651,38 @@ async def post_upload_rag_file( ) -> vertex_rag_data_service.UploadRagFileResponse: """Post-rpc interceptor for upload_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_upload_rag_file` interceptor runs + before the `post_upload_rag_file_with_metadata` interceptor. """ return response + async def post_upload_rag_file_with_metadata( + self, + response: vertex_rag_data_service.UploadRagFileResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.UploadRagFileResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for upload_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_upload_rag_file_with_metadata` + interceptor in new development instead of the `post_upload_rag_file` interceptor. + When both interceptors are used, this `post_upload_rag_file_with_metadata` interceptor runs after the + `post_upload_rag_file` interceptor. The (possibly modified) response returned by + `post_upload_rag_file` will be passed to + `post_upload_rag_file_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1002,6 +1241,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1156,6 +1399,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1310,6 +1557,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1464,6 +1715,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1619,6 +1874,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1779,6 +2038,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1932,6 +2195,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_rag_corpora(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_rag_corpora_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2087,6 +2354,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2249,6 +2520,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2408,6 +2683,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upload_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upload_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_service/client.py b/google/cloud/aiplatform_v1/services/vertex_rag_service/client.py index 4f8b6fe9f3..e51bc0d85d 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_service/client.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -503,6 +505,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1165,16 +1194,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1220,16 +1253,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1392,16 +1429,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1513,16 +1554,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1635,16 +1680,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1695,16 +1744,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1750,16 +1803,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1805,16 +1862,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest.py b/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest.py index 6f836581fa..22e754b987 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest.py @@ -127,12 +127,38 @@ def post_augment_prompt( ) -> vertex_rag_service.AugmentPromptResponse: """Post-rpc interceptor for augment_prompt - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_augment_prompt_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_augment_prompt` interceptor runs + before the `post_augment_prompt_with_metadata` interceptor. """ return response + def post_augment_prompt_with_metadata( + self, + response: vertex_rag_service.AugmentPromptResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.AugmentPromptResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for augment_prompt + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_augment_prompt_with_metadata` + interceptor in new development instead of the `post_augment_prompt` interceptor. + When both interceptors are used, this `post_augment_prompt_with_metadata` interceptor runs after the + `post_augment_prompt` interceptor. The (possibly modified) response returned by + `post_augment_prompt` will be passed to + `post_augment_prompt_with_metadata`. + """ + return response, metadata + def pre_corroborate_content( self, request: vertex_rag_service.CorroborateContentRequest, @@ -153,12 +179,38 @@ def post_corroborate_content( ) -> vertex_rag_service.CorroborateContentResponse: """Post-rpc interceptor for corroborate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_corroborate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_corroborate_content` interceptor runs + before the `post_corroborate_content_with_metadata` interceptor. """ return response + def post_corroborate_content_with_metadata( + self, + response: vertex_rag_service.CorroborateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.CorroborateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for corroborate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_corroborate_content_with_metadata` + interceptor in new development instead of the `post_corroborate_content` interceptor. + When both interceptors are used, this `post_corroborate_content_with_metadata` interceptor runs after the + `post_corroborate_content` interceptor. The (possibly modified) response returned by + `post_corroborate_content` will be passed to + `post_corroborate_content_with_metadata`. + """ + return response, metadata + def pre_retrieve_contexts( self, request: vertex_rag_service.RetrieveContextsRequest, @@ -179,12 +231,38 @@ def post_retrieve_contexts( ) -> vertex_rag_service.RetrieveContextsResponse: """Post-rpc interceptor for retrieve_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_retrieve_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_retrieve_contexts` interceptor runs + before the `post_retrieve_contexts_with_metadata` interceptor. """ return response + def post_retrieve_contexts_with_metadata( + self, + response: vertex_rag_service.RetrieveContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.RetrieveContextsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for retrieve_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_retrieve_contexts_with_metadata` + interceptor in new development instead of the `post_retrieve_contexts` interceptor. + When both interceptors are used, this `post_retrieve_contexts_with_metadata` interceptor runs after the + `post_retrieve_contexts` interceptor. The (possibly modified) response returned by + `post_retrieve_contexts` will be passed to + `post_retrieve_contexts_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -639,6 +717,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_augment_prompt(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_augment_prompt_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -794,6 +876,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_corroborate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_corroborate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -949,6 +1035,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_retrieve_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_retrieve_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest_asyncio.py index 7eca082979..ba20c8cc80 100644 --- a/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/vertex_rag_service/transports/rest_asyncio.py @@ -144,12 +144,38 @@ async def post_augment_prompt( ) -> vertex_rag_service.AugmentPromptResponse: """Post-rpc interceptor for augment_prompt - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_augment_prompt_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_augment_prompt` interceptor runs + before the `post_augment_prompt_with_metadata` interceptor. """ return response + async def post_augment_prompt_with_metadata( + self, + response: vertex_rag_service.AugmentPromptResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.AugmentPromptResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for augment_prompt + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_augment_prompt_with_metadata` + interceptor in new development instead of the `post_augment_prompt` interceptor. + When both interceptors are used, this `post_augment_prompt_with_metadata` interceptor runs after the + `post_augment_prompt` interceptor. The (possibly modified) response returned by + `post_augment_prompt` will be passed to + `post_augment_prompt_with_metadata`. + """ + return response, metadata + async def pre_corroborate_content( self, request: vertex_rag_service.CorroborateContentRequest, @@ -170,12 +196,38 @@ async def post_corroborate_content( ) -> vertex_rag_service.CorroborateContentResponse: """Post-rpc interceptor for corroborate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_corroborate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_corroborate_content` interceptor runs + before the `post_corroborate_content_with_metadata` interceptor. """ return response + async def post_corroborate_content_with_metadata( + self, + response: vertex_rag_service.CorroborateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.CorroborateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for corroborate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_corroborate_content_with_metadata` + interceptor in new development instead of the `post_corroborate_content` interceptor. + When both interceptors are used, this `post_corroborate_content_with_metadata` interceptor runs after the + `post_corroborate_content` interceptor. The (possibly modified) response returned by + `post_corroborate_content` will be passed to + `post_corroborate_content_with_metadata`. + """ + return response, metadata + async def pre_retrieve_contexts( self, request: vertex_rag_service.RetrieveContextsRequest, @@ -196,12 +248,38 @@ async def post_retrieve_contexts( ) -> vertex_rag_service.RetrieveContextsResponse: """Post-rpc interceptor for retrieve_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_retrieve_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_retrieve_contexts` interceptor runs + before the `post_retrieve_contexts_with_metadata` interceptor. """ return response + async def post_retrieve_contexts_with_metadata( + self, + response: vertex_rag_service.RetrieveContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.RetrieveContextsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for retrieve_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_retrieve_contexts_with_metadata` + interceptor in new development instead of the `post_retrieve_contexts` interceptor. + When both interceptors are used, this `post_retrieve_contexts_with_metadata` interceptor runs after the + `post_retrieve_contexts` interceptor. The (possibly modified) response returned by + `post_retrieve_contexts` will be passed to + `post_retrieve_contexts_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -720,6 +798,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_augment_prompt(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_augment_prompt_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -881,6 +963,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_corroborate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_corroborate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1042,6 +1128,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_retrieve_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_retrieve_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vizier_service/client.py b/google/cloud/aiplatform_v1/services/vizier_service/client.py index 03e830dfc2..8c2c5f52e1 100644 --- a/google/cloud/aiplatform_v1/services/vizier_service/client.py +++ b/google/cloud/aiplatform_v1/services/vizier_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -555,6 +557,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2407,16 +2436,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2462,16 +2495,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2634,16 +2671,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2755,16 +2796,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2877,16 +2922,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2937,16 +2986,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2992,16 +3045,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3047,16 +3104,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1/services/vizier_service/transports/rest.py b/google/cloud/aiplatform_v1/services/vizier_service/transports/rest.py index c08e170485..2a4c17ef08 100644 --- a/google/cloud/aiplatform_v1/services/vizier_service/transports/rest.py +++ b/google/cloud/aiplatform_v1/services/vizier_service/transports/rest.py @@ -218,12 +218,33 @@ def pre_add_trial_measurement( def post_add_trial_measurement(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for add_trial_measurement - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_trial_measurement_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_add_trial_measurement` interceptor runs + before the `post_add_trial_measurement_with_metadata` interceptor. """ return response + def post_add_trial_measurement_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for add_trial_measurement + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_add_trial_measurement_with_metadata` + interceptor in new development instead of the `post_add_trial_measurement` interceptor. + When both interceptors are used, this `post_add_trial_measurement_with_metadata` interceptor runs after the + `post_add_trial_measurement` interceptor. The (possibly modified) response returned by + `post_add_trial_measurement` will be passed to + `post_add_trial_measurement_with_metadata`. + """ + return response, metadata + def pre_check_trial_early_stopping_state( self, request: vizier_service.CheckTrialEarlyStoppingStateRequest, @@ -244,12 +265,35 @@ def post_check_trial_early_stopping_state( ) -> operations_pb2.Operation: """Post-rpc interceptor for check_trial_early_stopping_state - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_check_trial_early_stopping_state_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_check_trial_early_stopping_state` interceptor runs + before the `post_check_trial_early_stopping_state_with_metadata` interceptor. """ return response + def post_check_trial_early_stopping_state_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for check_trial_early_stopping_state + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_check_trial_early_stopping_state_with_metadata` + interceptor in new development instead of the `post_check_trial_early_stopping_state` interceptor. + When both interceptors are used, this `post_check_trial_early_stopping_state_with_metadata` interceptor runs after the + `post_check_trial_early_stopping_state` interceptor. The (possibly modified) response returned by + `post_check_trial_early_stopping_state` will be passed to + `post_check_trial_early_stopping_state_with_metadata`. + """ + return response, metadata + def pre_complete_trial( self, request: vizier_service.CompleteTrialRequest, @@ -267,12 +311,33 @@ def pre_complete_trial( def post_complete_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for complete_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_complete_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_complete_trial` interceptor runs + before the `post_complete_trial_with_metadata` interceptor. """ return response + def post_complete_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for complete_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_complete_trial_with_metadata` + interceptor in new development instead of the `post_complete_trial` interceptor. + When both interceptors are used, this `post_complete_trial_with_metadata` interceptor runs after the + `post_complete_trial` interceptor. The (possibly modified) response returned by + `post_complete_trial` will be passed to + `post_complete_trial_with_metadata`. + """ + return response, metadata + def pre_create_study( self, request: vizier_service.CreateStudyRequest, @@ -290,12 +355,35 @@ def pre_create_study( def post_create_study(self, response: gca_study.Study) -> gca_study.Study: """Post-rpc interceptor for create_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_study` interceptor runs + before the `post_create_study_with_metadata` interceptor. """ return response + def post_create_study_with_metadata( + self, + response: gca_study.Study, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_study_with_metadata` + interceptor in new development instead of the `post_create_study` interceptor. + When both interceptors are used, this `post_create_study_with_metadata` interceptor runs after the + `post_create_study` interceptor. The (possibly modified) response returned by + `post_create_study` will be passed to + `post_create_study_with_metadata`. + """ + return response, metadata + def pre_create_trial( self, request: vizier_service.CreateTrialRequest, @@ -313,12 +401,33 @@ def pre_create_trial( def post_create_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for create_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_trial` interceptor runs + before the `post_create_trial_with_metadata` interceptor. """ return response + def post_create_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_trial_with_metadata` + interceptor in new development instead of the `post_create_trial` interceptor. + When both interceptors are used, this `post_create_trial_with_metadata` interceptor runs after the + `post_create_trial` interceptor. The (possibly modified) response returned by + `post_create_trial` will be passed to + `post_create_trial_with_metadata`. + """ + return response, metadata + def pre_delete_study( self, request: vizier_service.DeleteStudyRequest, @@ -362,12 +471,33 @@ def pre_get_study( def post_get_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for get_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_study` interceptor runs + before the `post_get_study_with_metadata` interceptor. """ return response + def post_get_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_study_with_metadata` + interceptor in new development instead of the `post_get_study` interceptor. + When both interceptors are used, this `post_get_study_with_metadata` interceptor runs after the + `post_get_study` interceptor. The (possibly modified) response returned by + `post_get_study` will be passed to + `post_get_study_with_metadata`. + """ + return response, metadata + def pre_get_trial( self, request: vizier_service.GetTrialRequest, @@ -383,12 +513,33 @@ def pre_get_trial( def post_get_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for get_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_trial` interceptor runs + before the `post_get_trial_with_metadata` interceptor. """ return response + def post_get_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_trial_with_metadata` + interceptor in new development instead of the `post_get_trial` interceptor. + When both interceptors are used, this `post_get_trial_with_metadata` interceptor runs after the + `post_get_trial` interceptor. The (possibly modified) response returned by + `post_get_trial` will be passed to + `post_get_trial_with_metadata`. + """ + return response, metadata + def pre_list_optimal_trials( self, request: vizier_service.ListOptimalTrialsRequest, @@ -408,12 +559,38 @@ def post_list_optimal_trials( ) -> vizier_service.ListOptimalTrialsResponse: """Post-rpc interceptor for list_optimal_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_optimal_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_optimal_trials` interceptor runs + before the `post_list_optimal_trials_with_metadata` interceptor. """ return response + def post_list_optimal_trials_with_metadata( + self, + response: vizier_service.ListOptimalTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListOptimalTrialsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_optimal_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_optimal_trials_with_metadata` + interceptor in new development instead of the `post_list_optimal_trials` interceptor. + When both interceptors are used, this `post_list_optimal_trials_with_metadata` interceptor runs after the + `post_list_optimal_trials` interceptor. The (possibly modified) response returned by + `post_list_optimal_trials` will be passed to + `post_list_optimal_trials_with_metadata`. + """ + return response, metadata + def pre_list_studies( self, request: vizier_service.ListStudiesRequest, @@ -433,12 +610,37 @@ def post_list_studies( ) -> vizier_service.ListStudiesResponse: """Post-rpc interceptor for list_studies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_studies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_studies` interceptor runs + before the `post_list_studies_with_metadata` interceptor. """ return response + def post_list_studies_with_metadata( + self, + response: vizier_service.ListStudiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListStudiesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_studies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_studies_with_metadata` + interceptor in new development instead of the `post_list_studies` interceptor. + When both interceptors are used, this `post_list_studies_with_metadata` interceptor runs after the + `post_list_studies` interceptor. The (possibly modified) response returned by + `post_list_studies` will be passed to + `post_list_studies_with_metadata`. + """ + return response, metadata + def pre_list_trials( self, request: vizier_service.ListTrialsRequest, @@ -458,12 +660,37 @@ def post_list_trials( ) -> vizier_service.ListTrialsResponse: """Post-rpc interceptor for list_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_trials` interceptor runs + before the `post_list_trials_with_metadata` interceptor. """ return response + def post_list_trials_with_metadata( + self, + response: vizier_service.ListTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListTrialsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_trials_with_metadata` + interceptor in new development instead of the `post_list_trials` interceptor. + When both interceptors are used, this `post_list_trials_with_metadata` interceptor runs after the + `post_list_trials` interceptor. The (possibly modified) response returned by + `post_list_trials` will be passed to + `post_list_trials_with_metadata`. + """ + return response, metadata + def pre_lookup_study( self, request: vizier_service.LookupStudyRequest, @@ -481,12 +708,33 @@ def pre_lookup_study( def post_lookup_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for lookup_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_lookup_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_lookup_study` interceptor runs + before the `post_lookup_study_with_metadata` interceptor. """ return response + def post_lookup_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for lookup_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_lookup_study_with_metadata` + interceptor in new development instead of the `post_lookup_study` interceptor. + When both interceptors are used, this `post_lookup_study_with_metadata` interceptor runs after the + `post_lookup_study` interceptor. The (possibly modified) response returned by + `post_lookup_study` will be passed to + `post_lookup_study_with_metadata`. + """ + return response, metadata + def pre_stop_trial( self, request: vizier_service.StopTrialRequest, @@ -504,12 +752,33 @@ def pre_stop_trial( def post_stop_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for stop_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_stop_trial` interceptor runs + before the `post_stop_trial_with_metadata` interceptor. """ return response + def post_stop_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_stop_trial_with_metadata` + interceptor in new development instead of the `post_stop_trial` interceptor. + When both interceptors are used, this `post_stop_trial_with_metadata` interceptor runs after the + `post_stop_trial` interceptor. The (possibly modified) response returned by + `post_stop_trial` will be passed to + `post_stop_trial_with_metadata`. + """ + return response, metadata + def pre_suggest_trials( self, request: vizier_service.SuggestTrialsRequest, @@ -529,12 +798,35 @@ def post_suggest_trials( ) -> operations_pb2.Operation: """Post-rpc interceptor for suggest_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_suggest_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_suggest_trials` interceptor runs + before the `post_suggest_trials_with_metadata` interceptor. """ return response + def post_suggest_trials_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for suggest_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_suggest_trials_with_metadata` + interceptor in new development instead of the `post_suggest_trials` interceptor. + When both interceptors are used, this `post_suggest_trials_with_metadata` interceptor runs after the + `post_suggest_trials` interceptor. The (possibly modified) response returned by + `post_suggest_trials` will be passed to + `post_suggest_trials_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2916,6 +3208,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_trial_measurement(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_trial_measurement_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3071,6 +3367,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_check_trial_early_stopping_state(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_check_trial_early_stopping_state_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3224,6 +3527,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_complete_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_complete_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3376,6 +3683,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3533,6 +3844,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3899,6 +4214,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4050,6 +4369,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4202,6 +4525,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_optimal_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_optimal_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4350,6 +4677,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_studies(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_studies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4500,6 +4831,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4654,6 +4989,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_lookup_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_lookup_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4813,6 +5152,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_stop_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stop_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4962,6 +5305,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_suggest_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_suggest_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/services/vizier_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1/services/vizier_service/transports/rest_asyncio.py index d0c7b3b43d..22e1a4a5ee 100644 --- a/google/cloud/aiplatform_v1/services/vizier_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1/services/vizier_service/transports/rest_asyncio.py @@ -236,12 +236,33 @@ async def pre_add_trial_measurement( async def post_add_trial_measurement(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for add_trial_measurement - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_trial_measurement_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_add_trial_measurement` interceptor runs + before the `post_add_trial_measurement_with_metadata` interceptor. """ return response + async def post_add_trial_measurement_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for add_trial_measurement + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_add_trial_measurement_with_metadata` + interceptor in new development instead of the `post_add_trial_measurement` interceptor. + When both interceptors are used, this `post_add_trial_measurement_with_metadata` interceptor runs after the + `post_add_trial_measurement` interceptor. The (possibly modified) response returned by + `post_add_trial_measurement` will be passed to + `post_add_trial_measurement_with_metadata`. + """ + return response, metadata + async def pre_check_trial_early_stopping_state( self, request: vizier_service.CheckTrialEarlyStoppingStateRequest, @@ -262,12 +283,35 @@ async def post_check_trial_early_stopping_state( ) -> operations_pb2.Operation: """Post-rpc interceptor for check_trial_early_stopping_state - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_check_trial_early_stopping_state_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_check_trial_early_stopping_state` interceptor runs + before the `post_check_trial_early_stopping_state_with_metadata` interceptor. """ return response + async def post_check_trial_early_stopping_state_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for check_trial_early_stopping_state + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_check_trial_early_stopping_state_with_metadata` + interceptor in new development instead of the `post_check_trial_early_stopping_state` interceptor. + When both interceptors are used, this `post_check_trial_early_stopping_state_with_metadata` interceptor runs after the + `post_check_trial_early_stopping_state` interceptor. The (possibly modified) response returned by + `post_check_trial_early_stopping_state` will be passed to + `post_check_trial_early_stopping_state_with_metadata`. + """ + return response, metadata + async def pre_complete_trial( self, request: vizier_service.CompleteTrialRequest, @@ -285,12 +329,33 @@ async def pre_complete_trial( async def post_complete_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for complete_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_complete_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_complete_trial` interceptor runs + before the `post_complete_trial_with_metadata` interceptor. """ return response + async def post_complete_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for complete_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_complete_trial_with_metadata` + interceptor in new development instead of the `post_complete_trial` interceptor. + When both interceptors are used, this `post_complete_trial_with_metadata` interceptor runs after the + `post_complete_trial` interceptor. The (possibly modified) response returned by + `post_complete_trial` will be passed to + `post_complete_trial_with_metadata`. + """ + return response, metadata + async def pre_create_study( self, request: vizier_service.CreateStudyRequest, @@ -308,12 +373,35 @@ async def pre_create_study( async def post_create_study(self, response: gca_study.Study) -> gca_study.Study: """Post-rpc interceptor for create_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_study` interceptor runs + before the `post_create_study_with_metadata` interceptor. """ return response + async def post_create_study_with_metadata( + self, + response: gca_study.Study, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_study_with_metadata` + interceptor in new development instead of the `post_create_study` interceptor. + When both interceptors are used, this `post_create_study_with_metadata` interceptor runs after the + `post_create_study` interceptor. The (possibly modified) response returned by + `post_create_study` will be passed to + `post_create_study_with_metadata`. + """ + return response, metadata + async def pre_create_trial( self, request: vizier_service.CreateTrialRequest, @@ -331,12 +419,33 @@ async def pre_create_trial( async def post_create_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for create_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_trial` interceptor runs + before the `post_create_trial_with_metadata` interceptor. """ return response + async def post_create_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_trial_with_metadata` + interceptor in new development instead of the `post_create_trial` interceptor. + When both interceptors are used, this `post_create_trial_with_metadata` interceptor runs after the + `post_create_trial` interceptor. The (possibly modified) response returned by + `post_create_trial` will be passed to + `post_create_trial_with_metadata`. + """ + return response, metadata + async def pre_delete_study( self, request: vizier_service.DeleteStudyRequest, @@ -380,12 +489,33 @@ async def pre_get_study( async def post_get_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for get_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_study` interceptor runs + before the `post_get_study_with_metadata` interceptor. """ return response + async def post_get_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_study_with_metadata` + interceptor in new development instead of the `post_get_study` interceptor. + When both interceptors are used, this `post_get_study_with_metadata` interceptor runs after the + `post_get_study` interceptor. The (possibly modified) response returned by + `post_get_study` will be passed to + `post_get_study_with_metadata`. + """ + return response, metadata + async def pre_get_trial( self, request: vizier_service.GetTrialRequest, @@ -401,12 +531,33 @@ async def pre_get_trial( async def post_get_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for get_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_trial` interceptor runs + before the `post_get_trial_with_metadata` interceptor. """ return response + async def post_get_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_trial_with_metadata` + interceptor in new development instead of the `post_get_trial` interceptor. + When both interceptors are used, this `post_get_trial_with_metadata` interceptor runs after the + `post_get_trial` interceptor. The (possibly modified) response returned by + `post_get_trial` will be passed to + `post_get_trial_with_metadata`. + """ + return response, metadata + async def pre_list_optimal_trials( self, request: vizier_service.ListOptimalTrialsRequest, @@ -426,12 +577,38 @@ async def post_list_optimal_trials( ) -> vizier_service.ListOptimalTrialsResponse: """Post-rpc interceptor for list_optimal_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_optimal_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_optimal_trials` interceptor runs + before the `post_list_optimal_trials_with_metadata` interceptor. """ return response + async def post_list_optimal_trials_with_metadata( + self, + response: vizier_service.ListOptimalTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListOptimalTrialsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_optimal_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_optimal_trials_with_metadata` + interceptor in new development instead of the `post_list_optimal_trials` interceptor. + When both interceptors are used, this `post_list_optimal_trials_with_metadata` interceptor runs after the + `post_list_optimal_trials` interceptor. The (possibly modified) response returned by + `post_list_optimal_trials` will be passed to + `post_list_optimal_trials_with_metadata`. + """ + return response, metadata + async def pre_list_studies( self, request: vizier_service.ListStudiesRequest, @@ -451,12 +628,37 @@ async def post_list_studies( ) -> vizier_service.ListStudiesResponse: """Post-rpc interceptor for list_studies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_studies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_studies` interceptor runs + before the `post_list_studies_with_metadata` interceptor. """ return response + async def post_list_studies_with_metadata( + self, + response: vizier_service.ListStudiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListStudiesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_studies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_studies_with_metadata` + interceptor in new development instead of the `post_list_studies` interceptor. + When both interceptors are used, this `post_list_studies_with_metadata` interceptor runs after the + `post_list_studies` interceptor. The (possibly modified) response returned by + `post_list_studies` will be passed to + `post_list_studies_with_metadata`. + """ + return response, metadata + async def pre_list_trials( self, request: vizier_service.ListTrialsRequest, @@ -476,12 +678,37 @@ async def post_list_trials( ) -> vizier_service.ListTrialsResponse: """Post-rpc interceptor for list_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_trials` interceptor runs + before the `post_list_trials_with_metadata` interceptor. """ return response + async def post_list_trials_with_metadata( + self, + response: vizier_service.ListTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListTrialsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_trials_with_metadata` + interceptor in new development instead of the `post_list_trials` interceptor. + When both interceptors are used, this `post_list_trials_with_metadata` interceptor runs after the + `post_list_trials` interceptor. The (possibly modified) response returned by + `post_list_trials` will be passed to + `post_list_trials_with_metadata`. + """ + return response, metadata + async def pre_lookup_study( self, request: vizier_service.LookupStudyRequest, @@ -499,12 +726,33 @@ async def pre_lookup_study( async def post_lookup_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for lookup_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_lookup_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_lookup_study` interceptor runs + before the `post_lookup_study_with_metadata` interceptor. """ return response + async def post_lookup_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for lookup_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_lookup_study_with_metadata` + interceptor in new development instead of the `post_lookup_study` interceptor. + When both interceptors are used, this `post_lookup_study_with_metadata` interceptor runs after the + `post_lookup_study` interceptor. The (possibly modified) response returned by + `post_lookup_study` will be passed to + `post_lookup_study_with_metadata`. + """ + return response, metadata + async def pre_stop_trial( self, request: vizier_service.StopTrialRequest, @@ -522,12 +770,33 @@ async def pre_stop_trial( async def post_stop_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for stop_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_stop_trial` interceptor runs + before the `post_stop_trial_with_metadata` interceptor. """ return response + async def post_stop_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_stop_trial_with_metadata` + interceptor in new development instead of the `post_stop_trial` interceptor. + When both interceptors are used, this `post_stop_trial_with_metadata` interceptor runs after the + `post_stop_trial` interceptor. The (possibly modified) response returned by + `post_stop_trial` will be passed to + `post_stop_trial_with_metadata`. + """ + return response, metadata + async def pre_suggest_trials( self, request: vizier_service.SuggestTrialsRequest, @@ -547,12 +816,35 @@ async def post_suggest_trials( ) -> operations_pb2.Operation: """Post-rpc interceptor for suggest_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_suggest_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_suggest_trials` interceptor runs + before the `post_suggest_trials_with_metadata` interceptor. """ return response + async def post_suggest_trials_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for suggest_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_suggest_trials_with_metadata` + interceptor in new development instead of the `post_suggest_trials` interceptor. + When both interceptors are used, this `post_suggest_trials_with_metadata` interceptor runs after the + `post_suggest_trials` interceptor. The (possibly modified) response returned by + `post_suggest_trials` will be passed to + `post_suggest_trials_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1142,6 +1434,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_trial_measurement(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_trial_measurement_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1306,6 +1602,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_check_trial_early_stopping_state(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_check_trial_early_stopping_state_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1469,6 +1772,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_complete_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_complete_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1629,6 +1936,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1794,6 +2105,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2182,6 +2497,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2339,6 +2658,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2500,6 +2823,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_optimal_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_optimal_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2656,6 +2983,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_studies(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_studies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2814,6 +3145,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2976,6 +3311,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_lookup_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_lookup_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3143,6 +3482,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_stop_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stop_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3304,6 +3647,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_suggest_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_suggest_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1/types/__init__.py b/google/cloud/aiplatform_v1/types/__init__.py index 82bbf4d0ea..ee8c2ffbb4 100644 --- a/google/cloud/aiplatform_v1/types/__init__.py +++ b/google/cloud/aiplatform_v1/types/__init__.py @@ -744,9 +744,12 @@ ListModelEvaluationsResponse, ListModelsRequest, ListModelsResponse, + ListModelVersionCheckpointsRequest, + ListModelVersionCheckpointsResponse, ListModelVersionsRequest, ListModelVersionsResponse, MergeVersionAliasesRequest, + ModelVersionCheckpoint, UpdateExplanationDatasetOperationMetadata, UpdateExplanationDatasetRequest, UpdateExplanationDatasetResponse, @@ -813,6 +816,10 @@ UpgradeNotebookRuntimeResponse, NotebookExecutionJobView, ) +from .notebook_software_config import ( + NotebookSoftwareConfig, + PostStartupScriptConfig, +) from .openapi import ( Schema, Type, @@ -1743,9 +1750,12 @@ "ListModelEvaluationsResponse", "ListModelsRequest", "ListModelsResponse", + "ListModelVersionCheckpointsRequest", + "ListModelVersionCheckpointsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", "MergeVersionAliasesRequest", + "ModelVersionCheckpoint", "UpdateExplanationDatasetOperationMetadata", "UpdateExplanationDatasetRequest", "UpdateExplanationDatasetResponse", @@ -1795,6 +1805,8 @@ "UpgradeNotebookRuntimeRequest", "UpgradeNotebookRuntimeResponse", "NotebookExecutionJobView", + "NotebookSoftwareConfig", + "PostStartupScriptConfig", "Schema", "Type", "DeleteOperationMetadata", diff --git a/google/cloud/aiplatform_v1/types/feature_online_store_service.py b/google/cloud/aiplatform_v1/types/feature_online_store_service.py index f820dc5caa..bf112786f8 100644 --- a/google/cloud/aiplatform_v1/types/feature_online_store_service.py +++ b/google/cloud/aiplatform_v1/types/feature_online_store_service.py @@ -318,11 +318,11 @@ class StringFilter(proto.Message): class NumericFilter(proto.Message): r"""Numeric filter is used to search a subset of the entities by using boolean rules on numeric columns. For example: Database Point 0: - {name: “a” value_int: 42} {name: “b” value_float: 1.0} Database - Point 1: {name: “a” value_int: 10} {name: “b” value_float: 2.0} - Database Point 2: {name: “a” value_int: -1} {name: “b” value_float: - 3.0} Query: {name: “a” value_int: 12 operator: LESS} // Matches - Point 1, 2 {name: “b” value_float: 2.0 operator: EQUAL} // Matches + {name: "a" value_int: 42} {name: "b" value_float: 1.0} Database + Point 1: {name: "a" value_int: 10} {name: "b" value_float: 2.0} + Database Point 2: {name: "a" value_int: -1} {name: "b" value_float: + 3.0} Query: {name: "a" value_int: 12 operator: LESS} // Matches + Point 1, 2 {name: "b" value_float: 2.0 operator: EQUAL} // Matches Point 1 This message has `oneof`_ fields (mutually exclusive fields). @@ -356,7 +356,7 @@ class NumericFilter(proto.Message): """ class Operator(proto.Enum): - r"""Datapoints for which Operator is true relative to the query’s + r"""Datapoints for which Operator is true relative to the query's Value field will be allowlisted. Values: diff --git a/google/cloud/aiplatform_v1/types/feature_view.py b/google/cloud/aiplatform_v1/types/feature_view.py index 6e32f65586..cf296f4d10 100644 --- a/google/cloud/aiplatform_v1/types/feature_view.py +++ b/google/cloud/aiplatform_v1/types/feature_view.py @@ -19,6 +19,7 @@ import proto # type: ignore +from google.cloud.aiplatform_v1.types import machine_resources from google.protobuf import timestamp_pb2 # type: ignore @@ -100,6 +101,9 @@ class FeatureView(proto.Message): data, so that approximate nearest neighbor (a.k.a ANN) algorithms search can be performed during online serving. + optimized_config (google.cloud.aiplatform_v1.types.FeatureView.OptimizedConfig): + Optional. Configuration for FeatureView + created under Optimized FeatureOnlineStore. service_agent_type (google.cloud.aiplatform_v1.types.FeatureView.ServiceAgentType): Optional. Service agent type used during data sync. By default, the Vertex AI Service Agent is used. When using an @@ -413,6 +417,26 @@ class VertexRagSource(proto.Message): number=2, ) + class OptimizedConfig(proto.Message): + r"""Configuration for FeatureViews created in Optimized + FeatureOnlineStore. + + Attributes: + automatic_resources (google.cloud.aiplatform_v1.types.AutomaticResources): + Optional. A description of resources that the FeatureView + uses, which to large degree are decided by Vertex AI, and + optionally allows only a modest additional configuration. If + min_replica_count is not set, the default value is 2. If + max_replica_count is not set, the default value is 6. The + max allowed replica count is 1000. + """ + + automatic_resources: machine_resources.AutomaticResources = proto.Field( + proto.MESSAGE, + number=7, + message=machine_resources.AutomaticResources, + ) + big_query_source: BigQuerySource = proto.Field( proto.MESSAGE, number=6, @@ -464,6 +488,11 @@ class VertexRagSource(proto.Message): number=15, message=IndexConfig, ) + optimized_config: OptimizedConfig = proto.Field( + proto.MESSAGE, + number=16, + message=OptimizedConfig, + ) service_agent_type: ServiceAgentType = proto.Field( proto.ENUM, number=14, diff --git a/google/cloud/aiplatform_v1/types/model.py b/google/cloud/aiplatform_v1/types/model.py index de545a3fb6..63aa07029d 100644 --- a/google/cloud/aiplatform_v1/types/model.py +++ b/google/cloud/aiplatform_v1/types/model.py @@ -80,6 +80,8 @@ class Model(proto.Message): The description of the Model. version_description (str): The description of this version. + default_checkpoint_id (str): + The default checkpoint id of a model version. predict_schemata (google.cloud.aiplatform_v1.types.PredictSchemata): The schemata that describe formats of the Model's predictions and explanations as given and returned via @@ -556,6 +558,10 @@ class BaseModelSource(proto.Message): proto.STRING, number=30, ) + default_checkpoint_id: str = proto.Field( + proto.STRING, + number=53, + ) predict_schemata: "PredictSchemata" = proto.Field( proto.MESSAGE, number=4, @@ -1056,6 +1062,9 @@ class ModelContainerSpec(proto.Message): health_probe (google.cloud.aiplatform_v1.types.Probe): Immutable. Specification for Kubernetes readiness probe. + liveness_probe (google.cloud.aiplatform_v1.types.Probe): + Immutable. Specification for Kubernetes + liveness probe. """ image_uri: str = proto.Field( @@ -1112,6 +1121,11 @@ class ModelContainerSpec(proto.Message): number=13, message="Probe", ) + liveness_probe: "Probe" = proto.Field( + proto.MESSAGE, + number=14, + message="Probe", + ) class Port(proto.Message): @@ -1196,6 +1210,10 @@ class Probe(proto.Message): container to determine whether it is alive or ready to receive traffic. + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields @@ -1204,6 +1222,21 @@ class Probe(proto.Message): ExecAction probes the health of a container by executing a command. + This field is a member of `oneof`_ ``probe_type``. + http_get (google.cloud.aiplatform_v1.types.Probe.HttpGetAction): + HttpGetAction probes the health of a + container by sending an HTTP GET request. + + This field is a member of `oneof`_ ``probe_type``. + grpc (google.cloud.aiplatform_v1.types.Probe.GrpcAction): + GrpcAction probes the health of a container + by sending a gRPC request. + + This field is a member of `oneof`_ ``probe_type``. + tcp_socket (google.cloud.aiplatform_v1.types.Probe.TcpSocketAction): + TcpSocketAction probes the health of a + container by opening a TCP socket connection. + This field is a member of `oneof`_ ``probe_type``. period_seconds (int): How often (in seconds) to perform the probe. Default to 10 @@ -1217,6 +1250,26 @@ class Probe(proto.Message): period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'. + failure_threshold (int): + Number of consecutive failures before the + probe is considered failed. Defaults to 3. + Minimum value is 1. + + Maps to Kubernetes probe argument + 'failureThreshold'. + success_threshold (int): + Number of consecutive successes before the + probe is considered successful. Defaults to 1. + Minimum value is 1. + + Maps to Kubernetes probe argument + 'successThreshold'. + initial_delay_seconds (int): + Number of seconds to wait before starting the + probe. Defaults to 0. Minimum value is 0. + + Maps to Kubernetes probe argument + 'initialDelaySeconds'. """ class ExecAction(proto.Message): @@ -1238,12 +1291,147 @@ class ExecAction(proto.Message): number=1, ) + class HttpGetAction(proto.Message): + r"""HttpGetAction describes an action based on HTTP Get requests. + + Attributes: + path (str): + Path to access on the HTTP server. + port (int): + Number of the port to access on the + container. Number must be in the range 1 to + 65535. + host (str): + Host name to connect to, defaults to the + model serving container's IP. You probably want + to set "Host" in httpHeaders instead. + scheme (str): + Scheme to use for connecting to the host. + Defaults to HTTP. Acceptable values are "HTTP" + or "HTTPS". + http_headers (MutableSequence[google.cloud.aiplatform_v1.types.Probe.HttpHeader]): + Custom headers to set in the request. HTTP + allows repeated headers. + """ + + path: str = proto.Field( + proto.STRING, + number=1, + ) + port: int = proto.Field( + proto.INT32, + number=2, + ) + host: str = proto.Field( + proto.STRING, + number=3, + ) + scheme: str = proto.Field( + proto.STRING, + number=4, + ) + http_headers: MutableSequence["Probe.HttpHeader"] = proto.RepeatedField( + proto.MESSAGE, + number=5, + message="Probe.HttpHeader", + ) + + class GrpcAction(proto.Message): + r"""GrpcAction checks the health of a container using a gRPC + service. + + Attributes: + port (int): + Port number of the gRPC service. Number must + be in the range 1 to 65535. + service (str): + Service is the name of the service to place + in the gRPC HealthCheckRequest (see + https://github.com/grpc/grpc/blob/master/doc/health-checking.md). + + If this is not specified, the default behavior + is defined by gRPC. + """ + + port: int = proto.Field( + proto.INT32, + number=1, + ) + service: str = proto.Field( + proto.STRING, + number=2, + ) + + class TcpSocketAction(proto.Message): + r"""TcpSocketAction probes the health of a container by opening a + TCP socket connection. + + Attributes: + port (int): + Number of the port to access on the + container. Number must be in the range 1 to + 65535. + host (str): + Optional: Host name to connect to, defaults + to the model serving container's IP. + """ + + port: int = proto.Field( + proto.INT32, + number=1, + ) + host: str = proto.Field( + proto.STRING, + number=2, + ) + + class HttpHeader(proto.Message): + r"""HttpHeader describes a custom header to be used in HTTP + probes + + Attributes: + name (str): + The header field name. + This will be canonicalized upon output, so + case-variant names will be understood as the + same header. + value (str): + The header field value + """ + + name: str = proto.Field( + proto.STRING, + number=1, + ) + value: str = proto.Field( + proto.STRING, + number=2, + ) + exec_: ExecAction = proto.Field( proto.MESSAGE, number=1, oneof="probe_type", message=ExecAction, ) + http_get: HttpGetAction = proto.Field( + proto.MESSAGE, + number=4, + oneof="probe_type", + message=HttpGetAction, + ) + grpc: GrpcAction = proto.Field( + proto.MESSAGE, + number=5, + oneof="probe_type", + message=GrpcAction, + ) + tcp_socket: TcpSocketAction = proto.Field( + proto.MESSAGE, + number=6, + oneof="probe_type", + message=TcpSocketAction, + ) period_seconds: int = proto.Field( proto.INT32, number=2, @@ -1252,6 +1440,18 @@ class ExecAction(proto.Message): proto.INT32, number=3, ) + failure_threshold: int = proto.Field( + proto.INT32, + number=7, + ) + success_threshold: int = proto.Field( + proto.INT32, + number=8, + ) + initial_delay_seconds: int = proto.Field( + proto.INT32, + number=9, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/aiplatform_v1/types/model_service.py b/google/cloud/aiplatform_v1/types/model_service.py index 41a14d1fa6..d75e523b6c 100644 --- a/google/cloud/aiplatform_v1/types/model_service.py +++ b/google/cloud/aiplatform_v1/types/model_service.py @@ -41,6 +41,9 @@ "ListModelsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", + "ListModelVersionCheckpointsRequest", + "ModelVersionCheckpoint", + "ListModelVersionCheckpointsResponse", "UpdateModelRequest", "UpdateExplanationDatasetRequest", "UpdateExplanationDatasetOperationMetadata", @@ -403,6 +406,102 @@ def raw_page(self): ) +class ListModelVersionCheckpointsRequest(proto.Message): + r"""Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints]. + + Attributes: + name (str): + Required. The name of the model version to list checkpoints + for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest version + will be used. + page_size (int): + Optional. The standard list page size. + page_token (str): + Optional. The standard list page token. Typically obtained + via + [next_page_token][google.cloud.aiplatform.v1.ListModelVersionCheckpointsResponse.next_page_token] + of the previous + [ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + call. + """ + + name: str = proto.Field( + proto.STRING, + number=1, + ) + page_size: int = proto.Field( + proto.INT32, + number=2, + ) + page_token: str = proto.Field( + proto.STRING, + number=3, + ) + + +class ModelVersionCheckpoint(proto.Message): + r"""A proto representation of a Spanner-stored + ModelVersionCheckpoint. The meaning of the fields is equivalent + to their in-Spanner counterparts. + + Attributes: + checkpoint_id (str): + The ID of the checkpoint. + epoch (int): + The epoch of the checkpoint. + step (int): + The step of the checkpoint. + """ + + checkpoint_id: str = proto.Field( + proto.STRING, + number=1, + ) + epoch: int = proto.Field( + proto.INT64, + number=2, + ) + step: int = proto.Field( + proto.INT64, + number=3, + ) + + +class ListModelVersionCheckpointsResponse(proto.Message): + r"""Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints] + + Attributes: + checkpoints (MutableSequence[google.cloud.aiplatform_v1.types.ModelVersionCheckpoint]): + List of Model Version checkpoints. + next_page_token (str): + A token to retrieve the next page of results. Pass to + [ListModelVersionCheckpointsRequest.page_token][google.cloud.aiplatform.v1.ListModelVersionCheckpointsRequest.page_token] + to obtain that page. + """ + + @property + def raw_page(self): + return self + + checkpoints: MutableSequence["ModelVersionCheckpoint"] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message="ModelVersionCheckpoint", + ) + next_page_token: str = proto.Field( + proto.STRING, + number=2, + ) + + class UpdateModelRequest(proto.Message): r"""Request message for [ModelService.UpdateModel][google.cloud.aiplatform.v1.ModelService.UpdateModel]. diff --git a/google/cloud/aiplatform_v1/types/notebook_runtime.py b/google/cloud/aiplatform_v1/types/notebook_runtime.py index f930750fcc..18d6692f39 100644 --- a/google/cloud/aiplatform_v1/types/notebook_runtime.py +++ b/google/cloud/aiplatform_v1/types/notebook_runtime.py @@ -27,6 +27,7 @@ from google.cloud.aiplatform_v1.types import ( notebook_runtime_template_ref as gca_notebook_runtime_template_ref, ) +from google.cloud.aiplatform_v1.types import notebook_software_config from google.protobuf import timestamp_pb2 # type: ignore @@ -150,6 +151,9 @@ class NotebookRuntimeTemplate(proto.Message): encryption_spec (google.cloud.aiplatform_v1.types.EncryptionSpec): Customer-managed encryption key spec for the notebook runtime. + software_config (google.cloud.aiplatform_v1.types.NotebookSoftwareConfig): + Optional. The notebook software configuration + of the notebook runtime. """ name: str = proto.Field( @@ -237,6 +241,11 @@ class NotebookRuntimeTemplate(proto.Message): number=23, message=gca_encryption_spec.EncryptionSpec, ) + software_config: notebook_software_config.NotebookSoftwareConfig = proto.Field( + proto.MESSAGE, + number=24, + message=notebook_software_config.NotebookSoftwareConfig, + ) class NotebookRuntime(proto.Message): @@ -351,6 +360,9 @@ class NotebookRuntime(proto.Message): Optional. The Compute Engine tags to add to runtime (see `Tagging instances `__). + software_config (google.cloud.aiplatform_v1.types.NotebookSoftwareConfig): + Output only. Software config of the notebook + runtime. encryption_spec (google.cloud.aiplatform_v1.types.EncryptionSpec): Output only. Customer-managed encryption key spec for the notebook runtime. @@ -521,6 +533,11 @@ class RuntimeState(proto.Enum): proto.STRING, number=25, ) + software_config: notebook_software_config.NotebookSoftwareConfig = proto.Field( + proto.MESSAGE, + number=31, + message=notebook_software_config.NotebookSoftwareConfig, + ) encryption_spec: gca_encryption_spec.EncryptionSpec = proto.Field( proto.MESSAGE, number=28, diff --git a/google/cloud/aiplatform_v1/types/notebook_service.py b/google/cloud/aiplatform_v1/types/notebook_service.py index 39a6747dbd..a976ae1662 100644 --- a/google/cloud/aiplatform_v1/types/notebook_service.py +++ b/google/cloud/aiplatform_v1/types/notebook_service.py @@ -177,6 +177,8 @@ class ListNotebookRuntimeTemplatesRequest(proto.Message): - ``notebookRuntimeType`` supports = and !=. notebookRuntimeType enum: [USER_DEFINED, ONE_CLICK]. + - ``machineType`` supports = and !=. + - ``acceleratorType`` supports = and !=. Some examples: @@ -184,6 +186,8 @@ class ListNotebookRuntimeTemplatesRequest(proto.Message): - ``displayName="myDisplayName"`` - ``labels.myKey="myValue"`` - ``notebookRuntimeType=USER_DEFINED`` + - ``machineType=e2-standard-4`` + - ``acceleratorType=NVIDIA_TESLA_T4`` page_size (int): Optional. The standard list page size. page_token (str): @@ -438,6 +442,8 @@ class ListNotebookRuntimesRequest(proto.Message): UI_RESOURCE_STATE_CREATION_FAILED]. - ``notebookRuntimeType`` supports = and !=. notebookRuntimeType enum: [USER_DEFINED, ONE_CLICK]. + - ``machineType`` supports = and !=. + - ``acceleratorType`` supports = and !=. Some examples: @@ -450,6 +456,8 @@ class ListNotebookRuntimesRequest(proto.Message): - ``runtimeUser="test@google.com"`` - ``uiState=UI_RESOURCE_STATE_BEING_DELETED`` - ``notebookRuntimeType=USER_DEFINED`` + - ``machineType=e2-standard-4`` + - ``acceleratorType=NVIDIA_TESLA_T4`` page_size (int): Optional. The standard list page size. page_token (str): diff --git a/google/cloud/aiplatform_v1/types/notebook_software_config.py b/google/cloud/aiplatform_v1/types/notebook_software_config.py new file mode 100644 index 0000000000..380673a5c1 --- /dev/null +++ b/google/cloud/aiplatform_v1/types/notebook_software_config.py @@ -0,0 +1,108 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from typing import MutableMapping, MutableSequence + +import proto # type: ignore + +from google.cloud.aiplatform_v1.types import env_var + + +__protobuf__ = proto.module( + package="google.cloud.aiplatform.v1", + manifest={ + "PostStartupScriptConfig", + "NotebookSoftwareConfig", + }, +) + + +class PostStartupScriptConfig(proto.Message): + r"""Post startup script config. + + Attributes: + post_startup_script (str): + Optional. Post startup script to run after + runtime is started. + post_startup_script_url (str): + Optional. Post startup script url to + download. Example: https://bucket/script.sh + post_startup_script_behavior (google.cloud.aiplatform_v1.types.PostStartupScriptConfig.PostStartupScriptBehavior): + Optional. Post startup script behavior that + defines download and execution behavior. + """ + + class PostStartupScriptBehavior(proto.Enum): + r"""Represents a notebook runtime post startup script behavior. + + Values: + POST_STARTUP_SCRIPT_BEHAVIOR_UNSPECIFIED (0): + Unspecified post startup script behavior. + RUN_ONCE (1): + Run post startup script after runtime is + started. + RUN_EVERY_START (2): + Run post startup script after runtime is + stopped. + DOWNLOAD_AND_RUN_EVERY_START (3): + Download and run post startup script every + time runtime is started. + """ + POST_STARTUP_SCRIPT_BEHAVIOR_UNSPECIFIED = 0 + RUN_ONCE = 1 + RUN_EVERY_START = 2 + DOWNLOAD_AND_RUN_EVERY_START = 3 + + post_startup_script: str = proto.Field( + proto.STRING, + number=1, + ) + post_startup_script_url: str = proto.Field( + proto.STRING, + number=2, + ) + post_startup_script_behavior: PostStartupScriptBehavior = proto.Field( + proto.ENUM, + number=3, + enum=PostStartupScriptBehavior, + ) + + +class NotebookSoftwareConfig(proto.Message): + r"""Notebook Software Config. + + Attributes: + env (MutableSequence[google.cloud.aiplatform_v1.types.EnvVar]): + Optional. Environment variables to be passed + to the container. Maximum limit is 100. + post_startup_script_config (google.cloud.aiplatform_v1.types.PostStartupScriptConfig): + Optional. Post startup script config. + """ + + env: MutableSequence[env_var.EnvVar] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message=env_var.EnvVar, + ) + post_startup_script_config: "PostStartupScriptConfig" = proto.Field( + proto.MESSAGE, + number=2, + message="PostStartupScriptConfig", + ) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/aiplatform_v1/types/prediction_service.py b/google/cloud/aiplatform_v1/types/prediction_service.py index fd79b38f41..a1ce41cef8 100644 --- a/google/cloud/aiplatform_v1/types/prediction_service.py +++ b/google/cloud/aiplatform_v1/types/prediction_service.py @@ -25,6 +25,7 @@ from google.cloud.aiplatform_v1.types import tool from google.cloud.aiplatform_v1.types import types from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore __protobuf__ = proto.module( @@ -904,6 +905,12 @@ class GenerateContentResponse(proto.Message): model_version (str): Output only. The model version used to generate the response. + create_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. Timestamp when the request is + made to the server. + response_id (str): + Output only. response_id is used to identify each response. + It is the encoding of the event_id. prompt_feedback (google.cloud.aiplatform_v1.types.GenerateContentResponse.PromptFeedback): Output only. Content filter results for a prompt sent in the request. Note: Sent only in @@ -1039,6 +1046,15 @@ class UsageMetadata(proto.Message): proto.STRING, number=11, ) + create_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=12, + message=timestamp_pb2.Timestamp, + ) + response_id: str = proto.Field( + proto.STRING, + number=13, + ) prompt_feedback: PromptFeedback = proto.Field( proto.MESSAGE, number=3, diff --git a/google/cloud/aiplatform_v1beta1/__init__.py b/google/cloud/aiplatform_v1beta1/__init__.py index 5cc676def9..7ea733f9f4 100644 --- a/google/cloud/aiplatform_v1beta1/__init__.py +++ b/google/cloud/aiplatform_v1beta1/__init__.py @@ -211,6 +211,7 @@ from .types.endpoint import FasterDeploymentConfig from .types.endpoint import PredictRequestResponseLoggingConfig from .types.endpoint import PrivateEndpoints +from .types.endpoint import RolloutOptions from .types.endpoint_service import CreateEndpointOperationMetadata from .types.endpoint_service import CreateEndpointRequest from .types.endpoint_service import DeleteEndpointRequest @@ -248,8 +249,12 @@ from .types.evaluation_service import CometInstance from .types.evaluation_service import CometResult from .types.evaluation_service import CometSpec +from .types.evaluation_service import EvaluateDatasetOperationMetadata +from .types.evaluation_service import EvaluateDatasetRequest +from .types.evaluation_service import EvaluateDatasetResponse from .types.evaluation_service import EvaluateInstancesRequest from .types.evaluation_service import EvaluateInstancesResponse +from .types.evaluation_service import EvaluationDataset from .types.evaluation_service import ExactMatchInput from .types.evaluation_service import ExactMatchInstance from .types.evaluation_service import ExactMatchMetricValue @@ -267,10 +272,13 @@ from .types.evaluation_service import GroundednessInstance from .types.evaluation_service import GroundednessResult from .types.evaluation_service import GroundednessSpec +from .types.evaluation_service import Metric from .types.evaluation_service import MetricxInput from .types.evaluation_service import MetricxInstance from .types.evaluation_service import MetricxResult from .types.evaluation_service import MetricxSpec +from .types.evaluation_service import OutputConfig +from .types.evaluation_service import OutputInfo from .types.evaluation_service import PairwiseMetricInput from .types.evaluation_service import PairwiseMetricInstance from .types.evaluation_service import PairwiseMetricResult @@ -821,9 +829,12 @@ from .types.model_service import ListModelEvaluationsResponse from .types.model_service import ListModelsRequest from .types.model_service import ListModelsResponse +from .types.model_service import ListModelVersionCheckpointsRequest +from .types.model_service import ListModelVersionCheckpointsResponse from .types.model_service import ListModelVersionsRequest from .types.model_service import ListModelVersionsResponse from .types.model_service import MergeVersionAliasesRequest +from .types.model_service import ModelVersionCheckpoint from .types.model_service import UpdateExplanationDatasetOperationMetadata from .types.model_service import UpdateExplanationDatasetRequest from .types.model_service import UpdateExplanationDatasetResponse @@ -873,6 +884,8 @@ from .types.notebook_service import UpgradeNotebookRuntimeRequest from .types.notebook_service import UpgradeNotebookRuntimeResponse from .types.notebook_service import NotebookExecutionJobView +from .types.notebook_software_config import NotebookSoftwareConfig +from .types.notebook_software_config import PostStartupScriptConfig from .types.openapi import Schema from .types.openapi import Type from .types.operation import DeleteOperationMetadata @@ -1462,10 +1475,14 @@ "EntityType", "EnvVar", "ErrorAnalysisAnnotation", + "EvaluateDatasetOperationMetadata", + "EvaluateDatasetRequest", + "EvaluateDatasetResponse", "EvaluateInstancesRequest", "EvaluateInstancesResponse", "EvaluatedAnnotation", "EvaluatedAnnotationExplanation", + "EvaluationDataset", "EvaluationServiceClient", "Event", "ExactMatchInput", @@ -1726,6 +1743,8 @@ "ListModelMonitoringJobsResponse", "ListModelMonitorsRequest", "ListModelMonitorsResponse", + "ListModelVersionCheckpointsRequest", + "ListModelVersionCheckpointsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", "ListModelsRequest", @@ -1787,6 +1806,7 @@ "MetadataSchema", "MetadataServiceClient", "MetadataStore", + "Metric", "MetricxInput", "MetricxInstance", "MetricxResult", @@ -1831,6 +1851,7 @@ "ModelMonitoringTabularStats", "ModelServiceClient", "ModelSourceInfo", + "ModelVersionCheckpoint", "MutateDeployedIndexOperationMetadata", "MutateDeployedIndexRequest", "MutateDeployedIndexResponse", @@ -1857,6 +1878,9 @@ "NotebookRuntimeTemplateRef", "NotebookRuntimeType", "NotebookServiceClient", + "NotebookSoftwareConfig", + "OutputConfig", + "OutputInfo", "PSCAutomationConfig", "PairwiseChoice", "PairwiseMetricInput", @@ -1892,6 +1916,7 @@ "PointwiseMetricResult", "PointwiseMetricSpec", "Port", + "PostStartupScriptConfig", "PrebuiltVoiceConfig", "PredefinedSplit", "PredictLongRunningMetadata", @@ -1996,6 +2021,7 @@ "RetrievalMetadata", "RetrieveContextsRequest", "RetrieveContextsResponse", + "RolloutOptions", "RougeInput", "RougeInstance", "RougeMetricValue", diff --git a/google/cloud/aiplatform_v1beta1/gapic_metadata.json b/google/cloud/aiplatform_v1beta1/gapic_metadata.json index 59700b0786..84c1357ae3 100644 --- a/google/cloud/aiplatform_v1beta1/gapic_metadata.json +++ b/google/cloud/aiplatform_v1beta1/gapic_metadata.json @@ -577,6 +577,11 @@ "grpc": { "libraryClient": "EvaluationServiceClient", "rpcs": { + "EvaluateDataset": { + "methods": [ + "evaluate_dataset" + ] + }, "EvaluateInstances": { "methods": [ "evaluate_instances" @@ -587,6 +592,11 @@ "grpc-async": { "libraryClient": "EvaluationServiceAsyncClient", "rpcs": { + "EvaluateDataset": { + "methods": [ + "evaluate_dataset" + ] + }, "EvaluateInstances": { "methods": [ "evaluate_instances" @@ -597,6 +607,11 @@ "rest": { "libraryClient": "EvaluationServiceClient", "rpcs": { + "EvaluateDataset": { + "methods": [ + "evaluate_dataset" + ] + }, "EvaluateInstances": { "methods": [ "evaluate_instances" @@ -3668,6 +3683,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" @@ -3763,6 +3783,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" @@ -3858,6 +3883,11 @@ "list_model_evaluations" ] }, + "ListModelVersionCheckpoints": { + "methods": [ + "list_model_version_checkpoints" + ] + }, "ListModelVersions": { "methods": [ "list_model_versions" diff --git a/google/cloud/aiplatform_v1beta1/gapic_version.py b/google/cloud/aiplatform_v1beta1/gapic_version.py index 5c07cfea07..73cbb5f9a1 100644 --- a/google/cloud/aiplatform_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.80.0" # {x-release-please-version} +__version__ = "1.81.0" # {x-release-please-version} diff --git a/google/cloud/aiplatform_v1beta1/services/dataset_service/client.py b/google/cloud/aiplatform_v1beta1/services/dataset_service/client.py index 365df28aec..51534208c4 100644 --- a/google/cloud/aiplatform_v1beta1/services/dataset_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/dataset_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -641,6 +643,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3243,16 +3272,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3298,16 +3331,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3470,16 +3507,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3591,16 +3632,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3713,16 +3758,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3773,16 +3822,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3828,16 +3881,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3883,16 +3940,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest.py index 17848c7108..f151f71099 100644 --- a/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest.py @@ -261,12 +261,35 @@ def post_create_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset` interceptor runs + before the `post_create_dataset_with_metadata` interceptor. """ return response + def post_create_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_with_metadata` + interceptor in new development instead of the `post_create_dataset` interceptor. + When both interceptors are used, this `post_create_dataset_with_metadata` interceptor runs after the + `post_create_dataset` interceptor. The (possibly modified) response returned by + `post_create_dataset` will be passed to + `post_create_dataset_with_metadata`. + """ + return response, metadata + def pre_create_dataset_version( self, request: dataset_service.CreateDatasetVersionRequest, @@ -287,12 +310,35 @@ def post_create_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset_version` interceptor runs + before the `post_create_dataset_version_with_metadata` interceptor. """ return response + def post_create_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_version_with_metadata` + interceptor in new development instead of the `post_create_dataset_version` interceptor. + When both interceptors are used, this `post_create_dataset_version_with_metadata` interceptor runs after the + `post_create_dataset_version` interceptor. The (possibly modified) response returned by + `post_create_dataset_version` will be passed to + `post_create_dataset_version_with_metadata`. + """ + return response, metadata + def pre_delete_dataset( self, request: dataset_service.DeleteDatasetRequest, @@ -312,12 +358,35 @@ def post_delete_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset` interceptor runs + before the `post_delete_dataset_with_metadata` interceptor. """ return response + def post_delete_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_with_metadata` + interceptor in new development instead of the `post_delete_dataset` interceptor. + When both interceptors are used, this `post_delete_dataset_with_metadata` interceptor runs after the + `post_delete_dataset` interceptor. The (possibly modified) response returned by + `post_delete_dataset` will be passed to + `post_delete_dataset_with_metadata`. + """ + return response, metadata + def pre_delete_dataset_version( self, request: dataset_service.DeleteDatasetVersionRequest, @@ -338,12 +407,35 @@ def post_delete_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset_version` interceptor runs + before the `post_delete_dataset_version_with_metadata` interceptor. """ return response + def post_delete_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_version_with_metadata` + interceptor in new development instead of the `post_delete_dataset_version` interceptor. + When both interceptors are used, this `post_delete_dataset_version_with_metadata` interceptor runs after the + `post_delete_dataset_version` interceptor. The (possibly modified) response returned by + `post_delete_dataset_version` will be passed to + `post_delete_dataset_version_with_metadata`. + """ + return response, metadata + def pre_delete_saved_query( self, request: dataset_service.DeleteSavedQueryRequest, @@ -363,12 +455,35 @@ def post_delete_saved_query( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_saved_query - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_saved_query_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_saved_query` interceptor runs + before the `post_delete_saved_query_with_metadata` interceptor. """ return response + def post_delete_saved_query_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_saved_query + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_saved_query_with_metadata` + interceptor in new development instead of the `post_delete_saved_query` interceptor. + When both interceptors are used, this `post_delete_saved_query_with_metadata` interceptor runs after the + `post_delete_saved_query` interceptor. The (possibly modified) response returned by + `post_delete_saved_query` will be passed to + `post_delete_saved_query_with_metadata`. + """ + return response, metadata + def pre_export_data( self, request: dataset_service.ExportDataRequest, @@ -388,12 +503,35 @@ def post_export_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_export_data` interceptor runs + before the `post_export_data_with_metadata` interceptor. """ return response + def post_export_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_export_data_with_metadata` + interceptor in new development instead of the `post_export_data` interceptor. + When both interceptors are used, this `post_export_data_with_metadata` interceptor runs after the + `post_export_data` interceptor. The (possibly modified) response returned by + `post_export_data` will be passed to + `post_export_data_with_metadata`. + """ + return response, metadata + def pre_get_annotation_spec( self, request: dataset_service.GetAnnotationSpecRequest, @@ -414,12 +552,35 @@ def post_get_annotation_spec( ) -> annotation_spec.AnnotationSpec: """Post-rpc interceptor for get_annotation_spec - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_annotation_spec_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_annotation_spec` interceptor runs + before the `post_get_annotation_spec_with_metadata` interceptor. """ return response + def post_get_annotation_spec_with_metadata( + self, + response: annotation_spec.AnnotationSpec, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[annotation_spec.AnnotationSpec, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_annotation_spec + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_annotation_spec_with_metadata` + interceptor in new development instead of the `post_get_annotation_spec` interceptor. + When both interceptors are used, this `post_get_annotation_spec_with_metadata` interceptor runs after the + `post_get_annotation_spec` interceptor. The (possibly modified) response returned by + `post_get_annotation_spec` will be passed to + `post_get_annotation_spec_with_metadata`. + """ + return response, metadata + def pre_get_dataset( self, request: dataset_service.GetDatasetRequest, @@ -437,12 +598,35 @@ def pre_get_dataset( def post_get_dataset(self, response: dataset.Dataset) -> dataset.Dataset: """Post-rpc interceptor for get_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset` interceptor runs + before the `post_get_dataset_with_metadata` interceptor. """ return response + def post_get_dataset_with_metadata( + self, + response: dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_with_metadata` + interceptor in new development instead of the `post_get_dataset` interceptor. + When both interceptors are used, this `post_get_dataset_with_metadata` interceptor runs after the + `post_get_dataset` interceptor. The (possibly modified) response returned by + `post_get_dataset` will be passed to + `post_get_dataset_with_metadata`. + """ + return response, metadata + def pre_get_dataset_version( self, request: dataset_service.GetDatasetVersionRequest, @@ -463,12 +647,35 @@ def post_get_dataset_version( ) -> dataset_version.DatasetVersion: """Post-rpc interceptor for get_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset_version` interceptor runs + before the `post_get_dataset_version_with_metadata` interceptor. """ return response + def post_get_dataset_version_with_metadata( + self, + response: dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_version_with_metadata` + interceptor in new development instead of the `post_get_dataset_version` interceptor. + When both interceptors are used, this `post_get_dataset_version_with_metadata` interceptor runs after the + `post_get_dataset_version` interceptor. The (possibly modified) response returned by + `post_get_dataset_version` will be passed to + `post_get_dataset_version_with_metadata`. + """ + return response, metadata + def pre_import_data( self, request: dataset_service.ImportDataRequest, @@ -488,12 +695,35 @@ def post_import_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_import_data` interceptor runs + before the `post_import_data_with_metadata` interceptor. """ return response + def post_import_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_import_data_with_metadata` + interceptor in new development instead of the `post_import_data` interceptor. + When both interceptors are used, this `post_import_data_with_metadata` interceptor runs after the + `post_import_data` interceptor. The (possibly modified) response returned by + `post_import_data` will be passed to + `post_import_data_with_metadata`. + """ + return response, metadata + def pre_list_annotations( self, request: dataset_service.ListAnnotationsRequest, @@ -513,12 +743,37 @@ def post_list_annotations( ) -> dataset_service.ListAnnotationsResponse: """Post-rpc interceptor for list_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_annotations` interceptor runs + before the `post_list_annotations_with_metadata` interceptor. """ return response + def post_list_annotations_with_metadata( + self, + response: dataset_service.ListAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListAnnotationsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_annotations_with_metadata` + interceptor in new development instead of the `post_list_annotations` interceptor. + When both interceptors are used, this `post_list_annotations_with_metadata` interceptor runs after the + `post_list_annotations` interceptor. The (possibly modified) response returned by + `post_list_annotations` will be passed to + `post_list_annotations_with_metadata`. + """ + return response, metadata + def pre_list_data_items( self, request: dataset_service.ListDataItemsRequest, @@ -538,12 +793,37 @@ def post_list_data_items( ) -> dataset_service.ListDataItemsResponse: """Post-rpc interceptor for list_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_items` interceptor runs + before the `post_list_data_items_with_metadata` interceptor. """ return response + def post_list_data_items_with_metadata( + self, + response: dataset_service.ListDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_data_items_with_metadata` + interceptor in new development instead of the `post_list_data_items` interceptor. + When both interceptors are used, this `post_list_data_items_with_metadata` interceptor runs after the + `post_list_data_items` interceptor. The (possibly modified) response returned by + `post_list_data_items` will be passed to + `post_list_data_items_with_metadata`. + """ + return response, metadata + def pre_list_datasets( self, request: dataset_service.ListDatasetsRequest, @@ -563,12 +843,37 @@ def post_list_datasets( ) -> dataset_service.ListDatasetsResponse: """Post-rpc interceptor for list_datasets - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_datasets_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_datasets` interceptor runs + before the `post_list_datasets_with_metadata` interceptor. """ return response + def post_list_datasets_with_metadata( + self, + response: dataset_service.ListDatasetsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_datasets + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_datasets_with_metadata` + interceptor in new development instead of the `post_list_datasets` interceptor. + When both interceptors are used, this `post_list_datasets_with_metadata` interceptor runs after the + `post_list_datasets` interceptor. The (possibly modified) response returned by + `post_list_datasets` will be passed to + `post_list_datasets_with_metadata`. + """ + return response, metadata + def pre_list_dataset_versions( self, request: dataset_service.ListDatasetVersionsRequest, @@ -589,12 +894,38 @@ def post_list_dataset_versions( ) -> dataset_service.ListDatasetVersionsResponse: """Post-rpc interceptor for list_dataset_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_dataset_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_dataset_versions` interceptor runs + before the `post_list_dataset_versions_with_metadata` interceptor. """ return response + def post_list_dataset_versions_with_metadata( + self, + response: dataset_service.ListDatasetVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetVersionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_dataset_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_dataset_versions_with_metadata` + interceptor in new development instead of the `post_list_dataset_versions` interceptor. + When both interceptors are used, this `post_list_dataset_versions_with_metadata` interceptor runs after the + `post_list_dataset_versions` interceptor. The (possibly modified) response returned by + `post_list_dataset_versions` will be passed to + `post_list_dataset_versions_with_metadata`. + """ + return response, metadata + def pre_list_saved_queries( self, request: dataset_service.ListSavedQueriesRequest, @@ -614,12 +945,38 @@ def post_list_saved_queries( ) -> dataset_service.ListSavedQueriesResponse: """Post-rpc interceptor for list_saved_queries - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_saved_queries_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_saved_queries` interceptor runs + before the `post_list_saved_queries_with_metadata` interceptor. """ return response + def post_list_saved_queries_with_metadata( + self, + response: dataset_service.ListSavedQueriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListSavedQueriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_saved_queries + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_saved_queries_with_metadata` + interceptor in new development instead of the `post_list_saved_queries` interceptor. + When both interceptors are used, this `post_list_saved_queries_with_metadata` interceptor runs after the + `post_list_saved_queries` interceptor. The (possibly modified) response returned by + `post_list_saved_queries` will be passed to + `post_list_saved_queries_with_metadata`. + """ + return response, metadata + def pre_restore_dataset_version( self, request: dataset_service.RestoreDatasetVersionRequest, @@ -640,12 +997,35 @@ def post_restore_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for restore_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_restore_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_restore_dataset_version` interceptor runs + before the `post_restore_dataset_version_with_metadata` interceptor. """ return response + def post_restore_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for restore_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_restore_dataset_version_with_metadata` + interceptor in new development instead of the `post_restore_dataset_version` interceptor. + When both interceptors are used, this `post_restore_dataset_version_with_metadata` interceptor runs after the + `post_restore_dataset_version` interceptor. The (possibly modified) response returned by + `post_restore_dataset_version` will be passed to + `post_restore_dataset_version_with_metadata`. + """ + return response, metadata + def pre_search_data_items( self, request: dataset_service.SearchDataItemsRequest, @@ -665,12 +1045,37 @@ def post_search_data_items( ) -> dataset_service.SearchDataItemsResponse: """Post-rpc interceptor for search_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_search_data_items` interceptor runs + before the `post_search_data_items_with_metadata` interceptor. """ return response + def post_search_data_items_with_metadata( + self, + response: dataset_service.SearchDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.SearchDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for search_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_search_data_items_with_metadata` + interceptor in new development instead of the `post_search_data_items` interceptor. + When both interceptors are used, this `post_search_data_items_with_metadata` interceptor runs after the + `post_search_data_items` interceptor. The (possibly modified) response returned by + `post_search_data_items` will be passed to + `post_search_data_items_with_metadata`. + """ + return response, metadata + def pre_update_dataset( self, request: dataset_service.UpdateDatasetRequest, @@ -688,12 +1093,35 @@ def pre_update_dataset( def post_update_dataset(self, response: gca_dataset.Dataset) -> gca_dataset.Dataset: """Post-rpc interceptor for update_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset` interceptor runs + before the `post_update_dataset_with_metadata` interceptor. """ return response + def post_update_dataset_with_metadata( + self, + response: gca_dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_with_metadata` + interceptor in new development instead of the `post_update_dataset` interceptor. + When both interceptors are used, this `post_update_dataset_with_metadata` interceptor runs after the + `post_update_dataset` interceptor. The (possibly modified) response returned by + `post_update_dataset` will be passed to + `post_update_dataset_with_metadata`. + """ + return response, metadata + def pre_update_dataset_version( self, request: dataset_service.UpdateDatasetVersionRequest, @@ -714,12 +1142,37 @@ def post_update_dataset_version( ) -> gca_dataset_version.DatasetVersion: """Post-rpc interceptor for update_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset_version` interceptor runs + before the `post_update_dataset_version_with_metadata` interceptor. """ return response + def post_update_dataset_version_with_metadata( + self, + response: gca_dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_version_with_metadata` + interceptor in new development instead of the `post_update_dataset_version` interceptor. + When both interceptors are used, this `post_update_dataset_version_with_metadata` interceptor runs after the + `post_update_dataset_version` interceptor. The (possibly modified) response returned by + `post_update_dataset_version` will be passed to + `post_update_dataset_version_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3264,6 +3717,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3416,6 +3873,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3559,6 +4020,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3705,6 +4170,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3850,6 +4319,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_saved_query(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_saved_query_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4003,6 +4476,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4149,6 +4626,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_annotation_spec(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_annotation_spec_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4295,6 +4776,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4439,6 +4924,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4592,6 +5081,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4738,6 +5231,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4884,6 +5381,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5030,6 +5531,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_datasets(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_datasets_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5179,6 +5684,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_dataset_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_dataset_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5327,6 +5836,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_saved_queries(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_saved_queries_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5475,6 +5988,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_restore_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_restore_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5621,6 +6138,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5773,6 +6294,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5924,6 +6449,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest_asyncio.py index 1a691859d8..aa9004d022 100644 --- a/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/dataset_service/transports/rest_asyncio.py @@ -279,12 +279,35 @@ async def post_create_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset` interceptor runs + before the `post_create_dataset_with_metadata` interceptor. """ return response + async def post_create_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_with_metadata` + interceptor in new development instead of the `post_create_dataset` interceptor. + When both interceptors are used, this `post_create_dataset_with_metadata` interceptor runs after the + `post_create_dataset` interceptor. The (possibly modified) response returned by + `post_create_dataset` will be passed to + `post_create_dataset_with_metadata`. + """ + return response, metadata + async def pre_create_dataset_version( self, request: dataset_service.CreateDatasetVersionRequest, @@ -305,12 +328,35 @@ async def post_create_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_create_dataset_version` interceptor runs + before the `post_create_dataset_version_with_metadata` interceptor. """ return response + async def post_create_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_create_dataset_version_with_metadata` + interceptor in new development instead of the `post_create_dataset_version` interceptor. + When both interceptors are used, this `post_create_dataset_version_with_metadata` interceptor runs after the + `post_create_dataset_version` interceptor. The (possibly modified) response returned by + `post_create_dataset_version` will be passed to + `post_create_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_delete_dataset( self, request: dataset_service.DeleteDatasetRequest, @@ -330,12 +376,35 @@ async def post_delete_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset` interceptor runs + before the `post_delete_dataset_with_metadata` interceptor. """ return response + async def post_delete_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_with_metadata` + interceptor in new development instead of the `post_delete_dataset` interceptor. + When both interceptors are used, this `post_delete_dataset_with_metadata` interceptor runs after the + `post_delete_dataset` interceptor. The (possibly modified) response returned by + `post_delete_dataset` will be passed to + `post_delete_dataset_with_metadata`. + """ + return response, metadata + async def pre_delete_dataset_version( self, request: dataset_service.DeleteDatasetVersionRequest, @@ -356,12 +425,35 @@ async def post_delete_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_dataset_version` interceptor runs + before the `post_delete_dataset_version_with_metadata` interceptor. """ return response + async def post_delete_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_dataset_version_with_metadata` + interceptor in new development instead of the `post_delete_dataset_version` interceptor. + When both interceptors are used, this `post_delete_dataset_version_with_metadata` interceptor runs after the + `post_delete_dataset_version` interceptor. The (possibly modified) response returned by + `post_delete_dataset_version` will be passed to + `post_delete_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_delete_saved_query( self, request: dataset_service.DeleteSavedQueryRequest, @@ -381,12 +473,35 @@ async def post_delete_saved_query( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_saved_query - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_saved_query_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_delete_saved_query` interceptor runs + before the `post_delete_saved_query_with_metadata` interceptor. """ return response + async def post_delete_saved_query_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_saved_query + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_delete_saved_query_with_metadata` + interceptor in new development instead of the `post_delete_saved_query` interceptor. + When both interceptors are used, this `post_delete_saved_query_with_metadata` interceptor runs after the + `post_delete_saved_query` interceptor. The (possibly modified) response returned by + `post_delete_saved_query` will be passed to + `post_delete_saved_query_with_metadata`. + """ + return response, metadata + async def pre_export_data( self, request: dataset_service.ExportDataRequest, @@ -406,12 +521,35 @@ async def post_export_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_export_data` interceptor runs + before the `post_export_data_with_metadata` interceptor. """ return response + async def post_export_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_export_data_with_metadata` + interceptor in new development instead of the `post_export_data` interceptor. + When both interceptors are used, this `post_export_data_with_metadata` interceptor runs after the + `post_export_data` interceptor. The (possibly modified) response returned by + `post_export_data` will be passed to + `post_export_data_with_metadata`. + """ + return response, metadata + async def pre_get_annotation_spec( self, request: dataset_service.GetAnnotationSpecRequest, @@ -432,12 +570,35 @@ async def post_get_annotation_spec( ) -> annotation_spec.AnnotationSpec: """Post-rpc interceptor for get_annotation_spec - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_annotation_spec_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_annotation_spec` interceptor runs + before the `post_get_annotation_spec_with_metadata` interceptor. """ return response + async def post_get_annotation_spec_with_metadata( + self, + response: annotation_spec.AnnotationSpec, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[annotation_spec.AnnotationSpec, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_annotation_spec + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_annotation_spec_with_metadata` + interceptor in new development instead of the `post_get_annotation_spec` interceptor. + When both interceptors are used, this `post_get_annotation_spec_with_metadata` interceptor runs after the + `post_get_annotation_spec` interceptor. The (possibly modified) response returned by + `post_get_annotation_spec` will be passed to + `post_get_annotation_spec_with_metadata`. + """ + return response, metadata + async def pre_get_dataset( self, request: dataset_service.GetDatasetRequest, @@ -455,12 +616,35 @@ async def pre_get_dataset( async def post_get_dataset(self, response: dataset.Dataset) -> dataset.Dataset: """Post-rpc interceptor for get_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset` interceptor runs + before the `post_get_dataset_with_metadata` interceptor. """ return response + async def post_get_dataset_with_metadata( + self, + response: dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_with_metadata` + interceptor in new development instead of the `post_get_dataset` interceptor. + When both interceptors are used, this `post_get_dataset_with_metadata` interceptor runs after the + `post_get_dataset` interceptor. The (possibly modified) response returned by + `post_get_dataset` will be passed to + `post_get_dataset_with_metadata`. + """ + return response, metadata + async def pre_get_dataset_version( self, request: dataset_service.GetDatasetVersionRequest, @@ -481,12 +665,35 @@ async def post_get_dataset_version( ) -> dataset_version.DatasetVersion: """Post-rpc interceptor for get_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_get_dataset_version` interceptor runs + before the `post_get_dataset_version_with_metadata` interceptor. """ return response + async def post_get_dataset_version_with_metadata( + self, + response: dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_get_dataset_version_with_metadata` + interceptor in new development instead of the `post_get_dataset_version` interceptor. + When both interceptors are used, this `post_get_dataset_version_with_metadata` interceptor runs after the + `post_get_dataset_version` interceptor. The (possibly modified) response returned by + `post_get_dataset_version` will be passed to + `post_get_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_import_data( self, request: dataset_service.ImportDataRequest, @@ -506,12 +713,35 @@ async def post_import_data( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_import_data` interceptor runs + before the `post_import_data_with_metadata` interceptor. """ return response + async def post_import_data_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_import_data_with_metadata` + interceptor in new development instead of the `post_import_data` interceptor. + When both interceptors are used, this `post_import_data_with_metadata` interceptor runs after the + `post_import_data` interceptor. The (possibly modified) response returned by + `post_import_data` will be passed to + `post_import_data_with_metadata`. + """ + return response, metadata + async def pre_list_annotations( self, request: dataset_service.ListAnnotationsRequest, @@ -531,12 +761,37 @@ async def post_list_annotations( ) -> dataset_service.ListAnnotationsResponse: """Post-rpc interceptor for list_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_annotations` interceptor runs + before the `post_list_annotations_with_metadata` interceptor. """ return response + async def post_list_annotations_with_metadata( + self, + response: dataset_service.ListAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListAnnotationsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_annotations_with_metadata` + interceptor in new development instead of the `post_list_annotations` interceptor. + When both interceptors are used, this `post_list_annotations_with_metadata` interceptor runs after the + `post_list_annotations` interceptor. The (possibly modified) response returned by + `post_list_annotations` will be passed to + `post_list_annotations_with_metadata`. + """ + return response, metadata + async def pre_list_data_items( self, request: dataset_service.ListDataItemsRequest, @@ -556,12 +811,37 @@ async def post_list_data_items( ) -> dataset_service.ListDataItemsResponse: """Post-rpc interceptor for list_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_items` interceptor runs + before the `post_list_data_items_with_metadata` interceptor. """ return response + async def post_list_data_items_with_metadata( + self, + response: dataset_service.ListDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_data_items_with_metadata` + interceptor in new development instead of the `post_list_data_items` interceptor. + When both interceptors are used, this `post_list_data_items_with_metadata` interceptor runs after the + `post_list_data_items` interceptor. The (possibly modified) response returned by + `post_list_data_items` will be passed to + `post_list_data_items_with_metadata`. + """ + return response, metadata + async def pre_list_datasets( self, request: dataset_service.ListDatasetsRequest, @@ -581,12 +861,37 @@ async def post_list_datasets( ) -> dataset_service.ListDatasetsResponse: """Post-rpc interceptor for list_datasets - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_datasets_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_datasets` interceptor runs + before the `post_list_datasets_with_metadata` interceptor. """ return response + async def post_list_datasets_with_metadata( + self, + response: dataset_service.ListDatasetsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_datasets + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_datasets_with_metadata` + interceptor in new development instead of the `post_list_datasets` interceptor. + When both interceptors are used, this `post_list_datasets_with_metadata` interceptor runs after the + `post_list_datasets` interceptor. The (possibly modified) response returned by + `post_list_datasets` will be passed to + `post_list_datasets_with_metadata`. + """ + return response, metadata + async def pre_list_dataset_versions( self, request: dataset_service.ListDatasetVersionsRequest, @@ -607,12 +912,38 @@ async def post_list_dataset_versions( ) -> dataset_service.ListDatasetVersionsResponse: """Post-rpc interceptor for list_dataset_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_dataset_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_dataset_versions` interceptor runs + before the `post_list_dataset_versions_with_metadata` interceptor. """ return response + async def post_list_dataset_versions_with_metadata( + self, + response: dataset_service.ListDatasetVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListDatasetVersionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_dataset_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_dataset_versions_with_metadata` + interceptor in new development instead of the `post_list_dataset_versions` interceptor. + When both interceptors are used, this `post_list_dataset_versions_with_metadata` interceptor runs after the + `post_list_dataset_versions` interceptor. The (possibly modified) response returned by + `post_list_dataset_versions` will be passed to + `post_list_dataset_versions_with_metadata`. + """ + return response, metadata + async def pre_list_saved_queries( self, request: dataset_service.ListSavedQueriesRequest, @@ -632,12 +963,38 @@ async def post_list_saved_queries( ) -> dataset_service.ListSavedQueriesResponse: """Post-rpc interceptor for list_saved_queries - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_saved_queries_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_list_saved_queries` interceptor runs + before the `post_list_saved_queries_with_metadata` interceptor. """ return response + async def post_list_saved_queries_with_metadata( + self, + response: dataset_service.ListSavedQueriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.ListSavedQueriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_saved_queries + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_list_saved_queries_with_metadata` + interceptor in new development instead of the `post_list_saved_queries` interceptor. + When both interceptors are used, this `post_list_saved_queries_with_metadata` interceptor runs after the + `post_list_saved_queries` interceptor. The (possibly modified) response returned by + `post_list_saved_queries` will be passed to + `post_list_saved_queries_with_metadata`. + """ + return response, metadata + async def pre_restore_dataset_version( self, request: dataset_service.RestoreDatasetVersionRequest, @@ -658,12 +1015,35 @@ async def post_restore_dataset_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for restore_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_restore_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_restore_dataset_version` interceptor runs + before the `post_restore_dataset_version_with_metadata` interceptor. """ return response + async def post_restore_dataset_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for restore_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_restore_dataset_version_with_metadata` + interceptor in new development instead of the `post_restore_dataset_version` interceptor. + When both interceptors are used, this `post_restore_dataset_version_with_metadata` interceptor runs after the + `post_restore_dataset_version` interceptor. The (possibly modified) response returned by + `post_restore_dataset_version` will be passed to + `post_restore_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_search_data_items( self, request: dataset_service.SearchDataItemsRequest, @@ -683,12 +1063,37 @@ async def post_search_data_items( ) -> dataset_service.SearchDataItemsResponse: """Post-rpc interceptor for search_data_items - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_data_items_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_search_data_items` interceptor runs + before the `post_search_data_items_with_metadata` interceptor. """ return response + async def post_search_data_items_with_metadata( + self, + response: dataset_service.SearchDataItemsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + dataset_service.SearchDataItemsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for search_data_items + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_search_data_items_with_metadata` + interceptor in new development instead of the `post_search_data_items` interceptor. + When both interceptors are used, this `post_search_data_items_with_metadata` interceptor runs after the + `post_search_data_items` interceptor. The (possibly modified) response returned by + `post_search_data_items` will be passed to + `post_search_data_items_with_metadata`. + """ + return response, metadata + async def pre_update_dataset( self, request: dataset_service.UpdateDatasetRequest, @@ -708,12 +1113,35 @@ async def post_update_dataset( ) -> gca_dataset.Dataset: """Post-rpc interceptor for update_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset` interceptor runs + before the `post_update_dataset_with_metadata` interceptor. """ return response + async def post_update_dataset_with_metadata( + self, + response: gca_dataset.Dataset, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_dataset.Dataset, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_with_metadata` + interceptor in new development instead of the `post_update_dataset` interceptor. + When both interceptors are used, this `post_update_dataset_with_metadata` interceptor runs after the + `post_update_dataset` interceptor. The (possibly modified) response returned by + `post_update_dataset` will be passed to + `post_update_dataset_with_metadata`. + """ + return response, metadata + async def pre_update_dataset_version( self, request: dataset_service.UpdateDatasetVersionRequest, @@ -734,12 +1162,37 @@ async def post_update_dataset_version( ) -> gca_dataset_version.DatasetVersion: """Post-rpc interceptor for update_dataset_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_dataset_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DatasetService server but before - it is returned to user code. + it is returned to user code. This `post_update_dataset_version` interceptor runs + before the `post_update_dataset_version_with_metadata` interceptor. """ return response + async def post_update_dataset_version_with_metadata( + self, + response: gca_dataset_version.DatasetVersion, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_dataset_version.DatasetVersion, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_dataset_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DatasetService server but before it is returned to user code. + + We recommend only using this `post_update_dataset_version_with_metadata` + interceptor in new development instead of the `post_update_dataset_version` interceptor. + When both interceptors are used, this `post_update_dataset_version_with_metadata` interceptor runs after the + `post_update_dataset_version` interceptor. The (possibly modified) response returned by + `post_update_dataset_version` will be passed to + `post_update_dataset_version_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1345,6 +1798,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1505,6 +1962,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1660,6 +2121,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1814,6 +2279,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1970,6 +2439,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_saved_query(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_saved_query_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2133,6 +2606,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2288,6 +2765,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_annotation_spec(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_annotation_spec_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2442,6 +2923,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2595,6 +3080,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2758,6 +3247,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2913,6 +3406,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3069,6 +3566,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3225,6 +3726,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_datasets(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_datasets_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3380,6 +3885,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_dataset_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_dataset_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3537,6 +4046,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_saved_queries(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_saved_queries_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3693,6 +4206,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_restore_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_restore_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3848,6 +4368,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_data_items(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_search_data_items_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4010,6 +4534,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4167,6 +4695,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_dataset_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_dataset_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/client.py b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/client.py index f36a47e78a..f020a0dca2 100644 --- a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -590,6 +592,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1732,16 +1761,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1787,16 +1820,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1959,16 +1996,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2080,16 +2121,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2202,16 +2247,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2262,16 +2311,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2317,16 +2370,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2372,16 +2429,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest.py index 709885e6bd..faca6d6c28 100644 --- a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest.py @@ -154,12 +154,35 @@ def post_create_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_deployment_resource_pool` interceptor runs + before the `post_create_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_create_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_create_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_create_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_create_deployment_resource_pool_with_metadata` interceptor runs after the + `post_create_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_create_deployment_resource_pool` will be passed to + `post_create_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_delete_deployment_resource_pool( self, request: deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest, @@ -180,12 +203,35 @@ def post_delete_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_deployment_resource_pool` interceptor runs + before the `post_delete_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_delete_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_delete_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_delete_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_delete_deployment_resource_pool_with_metadata` interceptor runs after the + `post_delete_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_delete_deployment_resource_pool` will be passed to + `post_delete_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_get_deployment_resource_pool( self, request: deployment_resource_pool_service.GetDeploymentResourcePoolRequest, @@ -206,12 +252,38 @@ def post_get_deployment_resource_pool( ) -> deployment_resource_pool.DeploymentResourcePool: """Post-rpc interceptor for get_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_deployment_resource_pool` interceptor runs + before the `post_get_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_get_deployment_resource_pool_with_metadata( + self, + response: deployment_resource_pool.DeploymentResourcePool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool.DeploymentResourcePool, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_get_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_get_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_get_deployment_resource_pool_with_metadata` interceptor runs after the + `post_get_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_get_deployment_resource_pool` will be passed to + `post_get_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_list_deployment_resource_pools( self, request: deployment_resource_pool_service.ListDeploymentResourcePoolsRequest, @@ -233,12 +305,38 @@ def post_list_deployment_resource_pools( ) -> deployment_resource_pool_service.ListDeploymentResourcePoolsResponse: """Post-rpc interceptor for list_deployment_resource_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_deployment_resource_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_deployment_resource_pools` interceptor runs + before the `post_list_deployment_resource_pools_with_metadata` interceptor. """ return response + def post_list_deployment_resource_pools_with_metadata( + self, + response: deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_deployment_resource_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_list_deployment_resource_pools_with_metadata` + interceptor in new development instead of the `post_list_deployment_resource_pools` interceptor. + When both interceptors are used, this `post_list_deployment_resource_pools_with_metadata` interceptor runs after the + `post_list_deployment_resource_pools` interceptor. The (possibly modified) response returned by + `post_list_deployment_resource_pools` will be passed to + `post_list_deployment_resource_pools_with_metadata`. + """ + return response, metadata + def pre_query_deployed_models( self, request: deployment_resource_pool_service.QueryDeployedModelsRequest, @@ -259,12 +357,38 @@ def post_query_deployed_models( ) -> deployment_resource_pool_service.QueryDeployedModelsResponse: """Post-rpc interceptor for query_deployed_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_deployed_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_query_deployed_models` interceptor runs + before the `post_query_deployed_models_with_metadata` interceptor. """ return response + def post_query_deployed_models_with_metadata( + self, + response: deployment_resource_pool_service.QueryDeployedModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.QueryDeployedModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_deployed_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_query_deployed_models_with_metadata` + interceptor in new development instead of the `post_query_deployed_models` interceptor. + When both interceptors are used, this `post_query_deployed_models_with_metadata` interceptor runs after the + `post_query_deployed_models` interceptor. The (possibly modified) response returned by + `post_query_deployed_models` will be passed to + `post_query_deployed_models_with_metadata`. + """ + return response, metadata + def pre_update_deployment_resource_pool( self, request: deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest, @@ -285,12 +409,35 @@ def post_update_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_deployment_resource_pool` interceptor runs + before the `post_update_deployment_resource_pool_with_metadata` interceptor. """ return response + def post_update_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_update_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_update_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_update_deployment_resource_pool_with_metadata` interceptor runs after the + `post_update_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_update_deployment_resource_pool` will be passed to + `post_update_deployment_resource_pool_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2844,6 +2991,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2993,6 +3147,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3145,6 +3306,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3305,6 +3470,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_deployment_resource_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_deployment_resource_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3458,6 +3630,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_deployed_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_query_deployed_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3615,6 +3791,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest_asyncio.py index f3ec7ee9ca..cf68dea715 100644 --- a/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/deployment_resource_pool_service/transports/rest_asyncio.py @@ -172,12 +172,35 @@ async def post_create_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_deployment_resource_pool` interceptor runs + before the `post_create_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_create_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_create_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_create_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_create_deployment_resource_pool_with_metadata` interceptor runs after the + `post_create_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_create_deployment_resource_pool` will be passed to + `post_create_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_delete_deployment_resource_pool( self, request: deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest, @@ -198,12 +221,35 @@ async def post_delete_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_deployment_resource_pool` interceptor runs + before the `post_delete_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_delete_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_delete_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_delete_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_delete_deployment_resource_pool_with_metadata` interceptor runs after the + `post_delete_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_delete_deployment_resource_pool` will be passed to + `post_delete_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_get_deployment_resource_pool( self, request: deployment_resource_pool_service.GetDeploymentResourcePoolRequest, @@ -224,12 +270,38 @@ async def post_get_deployment_resource_pool( ) -> deployment_resource_pool.DeploymentResourcePool: """Post-rpc interceptor for get_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_deployment_resource_pool` interceptor runs + before the `post_get_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_get_deployment_resource_pool_with_metadata( + self, + response: deployment_resource_pool.DeploymentResourcePool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool.DeploymentResourcePool, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_get_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_get_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_get_deployment_resource_pool_with_metadata` interceptor runs after the + `post_get_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_get_deployment_resource_pool` will be passed to + `post_get_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_list_deployment_resource_pools( self, request: deployment_resource_pool_service.ListDeploymentResourcePoolsRequest, @@ -251,12 +323,38 @@ async def post_list_deployment_resource_pools( ) -> deployment_resource_pool_service.ListDeploymentResourcePoolsResponse: """Post-rpc interceptor for list_deployment_resource_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_deployment_resource_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_deployment_resource_pools` interceptor runs + before the `post_list_deployment_resource_pools_with_metadata` interceptor. """ return response + async def post_list_deployment_resource_pools_with_metadata( + self, + response: deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_deployment_resource_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_list_deployment_resource_pools_with_metadata` + interceptor in new development instead of the `post_list_deployment_resource_pools` interceptor. + When both interceptors are used, this `post_list_deployment_resource_pools_with_metadata` interceptor runs after the + `post_list_deployment_resource_pools` interceptor. The (possibly modified) response returned by + `post_list_deployment_resource_pools` will be passed to + `post_list_deployment_resource_pools_with_metadata`. + """ + return response, metadata + async def pre_query_deployed_models( self, request: deployment_resource_pool_service.QueryDeployedModelsRequest, @@ -277,12 +375,38 @@ async def post_query_deployed_models( ) -> deployment_resource_pool_service.QueryDeployedModelsResponse: """Post-rpc interceptor for query_deployed_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_deployed_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_query_deployed_models` interceptor runs + before the `post_query_deployed_models_with_metadata` interceptor. """ return response + async def post_query_deployed_models_with_metadata( + self, + response: deployment_resource_pool_service.QueryDeployedModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + deployment_resource_pool_service.QueryDeployedModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_deployed_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_query_deployed_models_with_metadata` + interceptor in new development instead of the `post_query_deployed_models` interceptor. + When both interceptors are used, this `post_query_deployed_models_with_metadata` interceptor runs after the + `post_query_deployed_models` interceptor. The (possibly modified) response returned by + `post_query_deployed_models` will be passed to + `post_query_deployed_models_with_metadata`. + """ + return response, metadata + async def pre_update_deployment_resource_pool( self, request: deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest, @@ -303,12 +427,35 @@ async def post_update_deployment_resource_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_deployment_resource_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_deployment_resource_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the DeploymentResourcePoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_deployment_resource_pool` interceptor runs + before the `post_update_deployment_resource_pool_with_metadata` interceptor. """ return response + async def post_update_deployment_resource_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_deployment_resource_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the DeploymentResourcePoolService server but before it is returned to user code. + + We recommend only using this `post_update_deployment_resource_pool_with_metadata` + interceptor in new development instead of the `post_update_deployment_resource_pool` interceptor. + When both interceptors are used, this `post_update_deployment_resource_pool_with_metadata` interceptor runs after the + `post_update_deployment_resource_pool` interceptor. The (possibly modified) response returned by + `post_update_deployment_resource_pool` will be passed to + `post_update_deployment_resource_pool_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -857,6 +1004,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1017,6 +1171,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1178,6 +1339,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1347,6 +1515,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_deployment_resource_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_deployment_resource_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1506,6 +1681,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_deployed_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_query_deployed_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1674,6 +1853,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_deployment_resource_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_deployment_resource_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/endpoint_service/client.py b/google/cloud/aiplatform_v1beta1/services/endpoint_service/client.py index 2437c6e0cf..38f57b84ff 100644 --- a/google/cloud/aiplatform_v1beta1/services/endpoint_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/endpoint_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -618,6 +620,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2137,16 +2166,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2192,16 +2225,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2364,16 +2401,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2485,16 +2526,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2607,16 +2652,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2667,16 +2716,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2722,16 +2775,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2777,16 +2834,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest.py index ad1f48bc98..6e9251ba3c 100644 --- a/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest.py @@ -178,12 +178,35 @@ def post_create_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_endpoint` interceptor runs + before the `post_create_endpoint_with_metadata` interceptor. """ return response + def post_create_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_create_endpoint_with_metadata` + interceptor in new development instead of the `post_create_endpoint` interceptor. + When both interceptors are used, this `post_create_endpoint_with_metadata` interceptor runs after the + `post_create_endpoint` interceptor. The (possibly modified) response returned by + `post_create_endpoint` will be passed to + `post_create_endpoint_with_metadata`. + """ + return response, metadata + def pre_delete_endpoint( self, request: endpoint_service.DeleteEndpointRequest, @@ -203,12 +226,35 @@ def post_delete_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_endpoint` interceptor runs + before the `post_delete_endpoint_with_metadata` interceptor. """ return response + def post_delete_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_endpoint` interceptor. + When both interceptors are used, this `post_delete_endpoint_with_metadata` interceptor runs after the + `post_delete_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_endpoint` will be passed to + `post_delete_endpoint_with_metadata`. + """ + return response, metadata + def pre_deploy_model( self, request: endpoint_service.DeployModelRequest, @@ -228,12 +274,35 @@ def post_deploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_model` interceptor runs + before the `post_deploy_model_with_metadata` interceptor. """ return response + def post_deploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_model_with_metadata` + interceptor in new development instead of the `post_deploy_model` interceptor. + When both interceptors are used, this `post_deploy_model_with_metadata` interceptor runs after the + `post_deploy_model` interceptor. The (possibly modified) response returned by + `post_deploy_model` will be passed to + `post_deploy_model_with_metadata`. + """ + return response, metadata + def pre_get_endpoint( self, request: endpoint_service.GetEndpointRequest, @@ -251,12 +320,35 @@ def pre_get_endpoint( def post_get_endpoint(self, response: endpoint.Endpoint) -> endpoint.Endpoint: """Post-rpc interceptor for get_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_endpoint` interceptor runs + before the `post_get_endpoint_with_metadata` interceptor. """ return response + def post_get_endpoint_with_metadata( + self, + response: endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_get_endpoint_with_metadata` + interceptor in new development instead of the `post_get_endpoint` interceptor. + When both interceptors are used, this `post_get_endpoint_with_metadata` interceptor runs after the + `post_get_endpoint` interceptor. The (possibly modified) response returned by + `post_get_endpoint` will be passed to + `post_get_endpoint_with_metadata`. + """ + return response, metadata + def pre_list_endpoints( self, request: endpoint_service.ListEndpointsRequest, @@ -276,12 +368,37 @@ def post_list_endpoints( ) -> endpoint_service.ListEndpointsResponse: """Post-rpc interceptor for list_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_endpoints` interceptor runs + before the `post_list_endpoints_with_metadata` interceptor. """ return response + def post_list_endpoints_with_metadata( + self, + response: endpoint_service.ListEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + endpoint_service.ListEndpointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_list_endpoints_with_metadata` + interceptor in new development instead of the `post_list_endpoints` interceptor. + When both interceptors are used, this `post_list_endpoints_with_metadata` interceptor runs after the + `post_list_endpoints` interceptor. The (possibly modified) response returned by + `post_list_endpoints` will be passed to + `post_list_endpoints_with_metadata`. + """ + return response, metadata + def pre_mutate_deployed_model( self, request: endpoint_service.MutateDeployedModelRequest, @@ -302,12 +419,35 @@ def post_mutate_deployed_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_model` interceptor runs + before the `post_mutate_deployed_model_with_metadata` interceptor. """ return response + def post_mutate_deployed_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_model_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_model` interceptor. + When both interceptors are used, this `post_mutate_deployed_model_with_metadata` interceptor runs after the + `post_mutate_deployed_model` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_model` will be passed to + `post_mutate_deployed_model_with_metadata`. + """ + return response, metadata + def pre_undeploy_model( self, request: endpoint_service.UndeployModelRequest, @@ -327,12 +467,35 @@ def post_undeploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_model` interceptor runs + before the `post_undeploy_model_with_metadata` interceptor. """ return response + def post_undeploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_model_with_metadata` + interceptor in new development instead of the `post_undeploy_model` interceptor. + When both interceptors are used, this `post_undeploy_model_with_metadata` interceptor runs after the + `post_undeploy_model` interceptor. The (possibly modified) response returned by + `post_undeploy_model` will be passed to + `post_undeploy_model_with_metadata`. + """ + return response, metadata + def pre_update_endpoint( self, request: endpoint_service.UpdateEndpointRequest, @@ -352,12 +515,35 @@ def post_update_endpoint( ) -> gca_endpoint.Endpoint: """Post-rpc interceptor for update_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint` interceptor runs + before the `post_update_endpoint_with_metadata` interceptor. """ return response + def post_update_endpoint_with_metadata( + self, + response: gca_endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_with_metadata` + interceptor in new development instead of the `post_update_endpoint` interceptor. + When both interceptors are used, this `post_update_endpoint_with_metadata` interceptor runs after the + `post_update_endpoint` interceptor. The (possibly modified) response returned by + `post_update_endpoint` will be passed to + `post_update_endpoint_with_metadata`. + """ + return response, metadata + def pre_update_endpoint_long_running( self, request: endpoint_service.UpdateEndpointLongRunningRequest, @@ -378,12 +564,35 @@ def post_update_endpoint_long_running( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_endpoint_long_running - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_long_running_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint_long_running` interceptor runs + before the `post_update_endpoint_long_running_with_metadata` interceptor. """ return response + def post_update_endpoint_long_running_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint_long_running + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_long_running_with_metadata` + interceptor in new development instead of the `post_update_endpoint_long_running` interceptor. + When both interceptors are used, this `post_update_endpoint_long_running_with_metadata` interceptor runs after the + `post_update_endpoint_long_running` interceptor. The (possibly modified) response returned by + `post_update_endpoint_long_running` will be passed to + `post_update_endpoint_long_running_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2927,6 +3136,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3070,6 +3283,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3219,6 +3436,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_deploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_deploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3364,6 +3585,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3508,6 +3733,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3662,6 +3891,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_mutate_deployed_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_mutate_deployed_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3811,6 +4044,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_undeploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_undeploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3962,6 +4199,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4117,6 +4358,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_endpoint_long_running(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_endpoint_long_running_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest_asyncio.py index f836976a52..dc5cafdbdb 100644 --- a/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/endpoint_service/transports/rest_asyncio.py @@ -196,12 +196,35 @@ async def post_create_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_endpoint` interceptor runs + before the `post_create_endpoint_with_metadata` interceptor. """ return response + async def post_create_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_create_endpoint_with_metadata` + interceptor in new development instead of the `post_create_endpoint` interceptor. + When both interceptors are used, this `post_create_endpoint_with_metadata` interceptor runs after the + `post_create_endpoint` interceptor. The (possibly modified) response returned by + `post_create_endpoint` will be passed to + `post_create_endpoint_with_metadata`. + """ + return response, metadata + async def pre_delete_endpoint( self, request: endpoint_service.DeleteEndpointRequest, @@ -221,12 +244,35 @@ async def post_delete_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_endpoint` interceptor runs + before the `post_delete_endpoint_with_metadata` interceptor. """ return response + async def post_delete_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_endpoint` interceptor. + When both interceptors are used, this `post_delete_endpoint_with_metadata` interceptor runs after the + `post_delete_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_endpoint` will be passed to + `post_delete_endpoint_with_metadata`. + """ + return response, metadata + async def pre_deploy_model( self, request: endpoint_service.DeployModelRequest, @@ -246,12 +292,35 @@ async def post_deploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_model` interceptor runs + before the `post_deploy_model_with_metadata` interceptor. """ return response + async def post_deploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_model_with_metadata` + interceptor in new development instead of the `post_deploy_model` interceptor. + When both interceptors are used, this `post_deploy_model_with_metadata` interceptor runs after the + `post_deploy_model` interceptor. The (possibly modified) response returned by + `post_deploy_model` will be passed to + `post_deploy_model_with_metadata`. + """ + return response, metadata + async def pre_get_endpoint( self, request: endpoint_service.GetEndpointRequest, @@ -269,12 +338,35 @@ async def pre_get_endpoint( async def post_get_endpoint(self, response: endpoint.Endpoint) -> endpoint.Endpoint: """Post-rpc interceptor for get_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_endpoint` interceptor runs + before the `post_get_endpoint_with_metadata` interceptor. """ return response + async def post_get_endpoint_with_metadata( + self, + response: endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_get_endpoint_with_metadata` + interceptor in new development instead of the `post_get_endpoint` interceptor. + When both interceptors are used, this `post_get_endpoint_with_metadata` interceptor runs after the + `post_get_endpoint` interceptor. The (possibly modified) response returned by + `post_get_endpoint` will be passed to + `post_get_endpoint_with_metadata`. + """ + return response, metadata + async def pre_list_endpoints( self, request: endpoint_service.ListEndpointsRequest, @@ -294,12 +386,37 @@ async def post_list_endpoints( ) -> endpoint_service.ListEndpointsResponse: """Post-rpc interceptor for list_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_endpoints` interceptor runs + before the `post_list_endpoints_with_metadata` interceptor. """ return response + async def post_list_endpoints_with_metadata( + self, + response: endpoint_service.ListEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + endpoint_service.ListEndpointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_list_endpoints_with_metadata` + interceptor in new development instead of the `post_list_endpoints` interceptor. + When both interceptors are used, this `post_list_endpoints_with_metadata` interceptor runs after the + `post_list_endpoints` interceptor. The (possibly modified) response returned by + `post_list_endpoints` will be passed to + `post_list_endpoints_with_metadata`. + """ + return response, metadata + async def pre_mutate_deployed_model( self, request: endpoint_service.MutateDeployedModelRequest, @@ -320,12 +437,35 @@ async def post_mutate_deployed_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_model` interceptor runs + before the `post_mutate_deployed_model_with_metadata` interceptor. """ return response + async def post_mutate_deployed_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_model_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_model` interceptor. + When both interceptors are used, this `post_mutate_deployed_model_with_metadata` interceptor runs after the + `post_mutate_deployed_model` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_model` will be passed to + `post_mutate_deployed_model_with_metadata`. + """ + return response, metadata + async def pre_undeploy_model( self, request: endpoint_service.UndeployModelRequest, @@ -345,12 +485,35 @@ async def post_undeploy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_model` interceptor runs + before the `post_undeploy_model_with_metadata` interceptor. """ return response + async def post_undeploy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_model_with_metadata` + interceptor in new development instead of the `post_undeploy_model` interceptor. + When both interceptors are used, this `post_undeploy_model_with_metadata` interceptor runs after the + `post_undeploy_model` interceptor. The (possibly modified) response returned by + `post_undeploy_model` will be passed to + `post_undeploy_model_with_metadata`. + """ + return response, metadata + async def pre_update_endpoint( self, request: endpoint_service.UpdateEndpointRequest, @@ -370,12 +533,35 @@ async def post_update_endpoint( ) -> gca_endpoint.Endpoint: """Post-rpc interceptor for update_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint` interceptor runs + before the `post_update_endpoint_with_metadata` interceptor. """ return response + async def post_update_endpoint_with_metadata( + self, + response: gca_endpoint.Endpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_endpoint.Endpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_with_metadata` + interceptor in new development instead of the `post_update_endpoint` interceptor. + When both interceptors are used, this `post_update_endpoint_with_metadata` interceptor runs after the + `post_update_endpoint` interceptor. The (possibly modified) response returned by + `post_update_endpoint` will be passed to + `post_update_endpoint_with_metadata`. + """ + return response, metadata + async def pre_update_endpoint_long_running( self, request: endpoint_service.UpdateEndpointLongRunningRequest, @@ -396,12 +582,35 @@ async def post_update_endpoint_long_running( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_endpoint_long_running - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_endpoint_long_running_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_endpoint_long_running` interceptor runs + before the `post_update_endpoint_long_running_with_metadata` interceptor. """ return response + async def post_update_endpoint_long_running_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_endpoint_long_running + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EndpointService server but before it is returned to user code. + + We recommend only using this `post_update_endpoint_long_running_with_metadata` + interceptor in new development instead of the `post_update_endpoint_long_running` interceptor. + When both interceptors are used, this `post_update_endpoint_long_running_with_metadata` interceptor runs after the + `post_update_endpoint_long_running` interceptor. The (possibly modified) response returned by + `post_update_endpoint_long_running` will be passed to + `post_update_endpoint_long_running_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -957,6 +1166,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1113,6 +1326,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1274,6 +1491,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_deploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_deploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1429,6 +1650,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1584,6 +1809,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1746,6 +1975,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_mutate_deployed_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_mutate_deployed_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1908,6 +2141,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_undeploy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_undeploy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2070,6 +2307,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2234,6 +2475,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_endpoint_long_running(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_endpoint_long_running_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/async_client.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/async_client.py index 4c76357e3e..c99d451522 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/async_client.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/async_client.py @@ -44,6 +44,8 @@ except AttributeError: # pragma: NO COVER OptionalRetry = Union[retries.AsyncRetry, object, None] # type: ignore +from google.api_core import operation # type: ignore +from google.api_core import operation_async # type: ignore from google.cloud.aiplatform_v1beta1.types import evaluation_service from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore @@ -379,6 +381,117 @@ async def sample_evaluate_instances(): # Done; return the response. return response + async def evaluate_dataset( + self, + request: Optional[ + Union[evaluation_service.EvaluateDatasetRequest, dict] + ] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> operation_async.AsyncOperation: + r"""Evaluates a dataset based on a set of given metrics. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1beta1 + + async def sample_evaluate_dataset(): + # Create a client + client = aiplatform_v1beta1.EvaluationServiceAsyncClient() + + # Initialize request argument(s) + dataset = aiplatform_v1beta1.EvaluationDataset() + dataset.gcs_source.uris = ['uris_value1', 'uris_value2'] + + output_config = aiplatform_v1beta1.OutputConfig() + output_config.gcs_destination.output_uri_prefix = "output_uri_prefix_value" + + request = aiplatform_v1beta1.EvaluateDatasetRequest( + location="location_value", + dataset=dataset, + output_config=output_config, + ) + + # Make the request + operation = client.evaluate_dataset(request=request) + + print("Waiting for operation to complete...") + + response = (await operation).result() + + # Handle the response + print(response) + + Args: + request (Optional[Union[google.cloud.aiplatform_v1beta1.types.EvaluateDatasetRequest, dict]]): + The request object. Request message for + EvaluationService.EvaluateDataset. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.api_core.operation_async.AsyncOperation: + An object representing a long-running operation. + + The result type for the operation will be + :class:`google.cloud.aiplatform_v1beta1.types.EvaluateDatasetResponse` + Response in LRO for EvaluationService.EvaluateDataset. + + """ + # Create or coerce a protobuf request object. + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, evaluation_service.EvaluateDatasetRequest): + request = evaluation_service.EvaluateDatasetRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._client._transport._wrapped_methods[ + self._client._transport.evaluate_dataset + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("location", request.location),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Wrap the response in an operation future. + response = operation_async.from_gapic( + response, + self._client._transport.operations_client, + evaluation_service.EvaluateDatasetResponse, + metadata_type=evaluation_service.EvaluateDatasetOperationMetadata, + ) + + # Done; return the response. + return response + async def list_operations( self, request: Optional[operations_pb2.ListOperationsRequest] = None, diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/client.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/client.py index 63cdb7435a..13ee441df1 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -58,6 +60,8 @@ _LOGGER = std_logging.getLogger(__name__) +from google.api_core import operation # type: ignore +from google.api_core import operation_async # type: ignore from google.cloud.aiplatform_v1beta1.types import evaluation_service from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore @@ -478,6 +482,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -790,6 +821,115 @@ def sample_evaluate_instances(): # Done; return the response. return response + def evaluate_dataset( + self, + request: Optional[ + Union[evaluation_service.EvaluateDatasetRequest, dict] + ] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> operation.Operation: + r"""Evaluates a dataset based on a set of given metrics. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1beta1 + + def sample_evaluate_dataset(): + # Create a client + client = aiplatform_v1beta1.EvaluationServiceClient() + + # Initialize request argument(s) + dataset = aiplatform_v1beta1.EvaluationDataset() + dataset.gcs_source.uris = ['uris_value1', 'uris_value2'] + + output_config = aiplatform_v1beta1.OutputConfig() + output_config.gcs_destination.output_uri_prefix = "output_uri_prefix_value" + + request = aiplatform_v1beta1.EvaluateDatasetRequest( + location="location_value", + dataset=dataset, + output_config=output_config, + ) + + # Make the request + operation = client.evaluate_dataset(request=request) + + print("Waiting for operation to complete...") + + response = operation.result() + + # Handle the response + print(response) + + Args: + request (Union[google.cloud.aiplatform_v1beta1.types.EvaluateDatasetRequest, dict]): + The request object. Request message for + EvaluationService.EvaluateDataset. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.api_core.operation.Operation: + An object representing a long-running operation. + + The result type for the operation will be + :class:`google.cloud.aiplatform_v1beta1.types.EvaluateDatasetResponse` + Response in LRO for EvaluationService.EvaluateDataset. + + """ + # Create or coerce a protobuf request object. + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, evaluation_service.EvaluateDatasetRequest): + request = evaluation_service.EvaluateDatasetRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.evaluate_dataset] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("location", request.location),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Wrap the response in an operation future. + response = operation.from_gapic( + response, + self._transport.operations_client, + evaluation_service.EvaluateDatasetResponse, + metadata_type=evaluation_service.EvaluateDatasetOperationMetadata, + ) + + # Done; return the response. + return response + def __enter__(self) -> "EvaluationServiceClient": return self @@ -847,16 +987,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -902,16 +1046,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1074,16 +1222,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1195,16 +1347,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1317,16 +1473,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1377,16 +1537,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1432,16 +1596,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1487,16 +1655,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/base.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/base.py index ad98218e1e..9de6793bb4 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/base.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/base.py @@ -23,6 +23,7 @@ from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 from google.api_core import retry as retries +from google.api_core import operations_v1 from google.auth import credentials as ga_credentials # type: ignore from google.oauth2 import service_account # type: ignore @@ -138,6 +139,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=60.0, client_info=client_info, ), + self.evaluate_dataset: gapic_v1.method.wrap_method( + self.evaluate_dataset, + default_timeout=None, + client_info=client_info, + ), self.get_location: gapic_v1.method.wrap_method( self.get_location, default_timeout=None, @@ -199,6 +205,11 @@ def close(self): """ raise NotImplementedError() + @property + def operations_client(self): + """Return the client designed to process long-running operations.""" + raise NotImplementedError() + @property def evaluate_instances( self, @@ -211,6 +222,15 @@ def evaluate_instances( ]: raise NotImplementedError() + @property + def evaluate_dataset( + self, + ) -> Callable[ + [evaluation_service.EvaluateDatasetRequest], + Union[operations_pb2.Operation, Awaitable[operations_pb2.Operation]], + ]: + raise NotImplementedError() + @property def list_operations( self, diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc.py index cb04eaab97..8130497863 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc.py @@ -20,6 +20,7 @@ from typing import Callable, Dict, Optional, Sequence, Tuple, Union from google.api_core import grpc_helpers +from google.api_core import operations_v1 from google.api_core import gapic_v1 import google.auth # type: ignore from google.auth import credentials as ga_credentials # type: ignore @@ -199,6 +200,7 @@ def __init__( self._grpc_channel = None self._ssl_channel_credentials = ssl_channel_credentials self._stubs: Dict[str, Callable] = {} + self._operations_client: Optional[operations_v1.OperationsClient] = None if api_mtls_endpoint: warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning) @@ -325,6 +327,22 @@ def grpc_channel(self) -> grpc.Channel: """Return the channel designed to connect to this service.""" return self._grpc_channel + @property + def operations_client(self) -> operations_v1.OperationsClient: + """Create the client designed to process long-running operations. + + This property caches on the instance; repeated calls return the same + client. + """ + # Quick check: Only create a new client if we do not already have one. + if self._operations_client is None: + self._operations_client = operations_v1.OperationsClient( + self._logged_channel + ) + + # Return the client from cache. + return self._operations_client + @property def evaluate_instances( self, @@ -354,6 +372,34 @@ def evaluate_instances( ) return self._stubs["evaluate_instances"] + @property + def evaluate_dataset( + self, + ) -> Callable[ + [evaluation_service.EvaluateDatasetRequest], operations_pb2.Operation + ]: + r"""Return a callable for the evaluate dataset method over gRPC. + + Evaluates a dataset based on a set of given metrics. + + Returns: + Callable[[~.EvaluateDatasetRequest], + ~.Operation]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "evaluate_dataset" not in self._stubs: + self._stubs["evaluate_dataset"] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1beta1.EvaluationService/EvaluateDataset", + request_serializer=evaluation_service.EvaluateDatasetRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["evaluate_dataset"] + def close(self): self._logged_channel.close() diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc_asyncio.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc_asyncio.py index 0ef1ec438a..1fb393c9db 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/grpc_asyncio.py @@ -24,6 +24,7 @@ from google.api_core import grpc_helpers_async from google.api_core import exceptions as core_exceptions from google.api_core import retry_async as retries +from google.api_core import operations_v1 from google.auth import credentials as ga_credentials # type: ignore from google.auth.transport.grpc import SslCredentials # type: ignore from google.protobuf.json_format import MessageToJson @@ -248,6 +249,7 @@ def __init__( self._grpc_channel = None self._ssl_channel_credentials = ssl_channel_credentials self._stubs: Dict[str, Callable] = {} + self._operations_client: Optional[operations_v1.OperationsAsyncClient] = None if api_mtls_endpoint: warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning) @@ -332,6 +334,22 @@ def grpc_channel(self) -> aio.Channel: # Return the channel from cache. return self._grpc_channel + @property + def operations_client(self) -> operations_v1.OperationsAsyncClient: + """Create the client designed to process long-running operations. + + This property caches on the instance; repeated calls return the same + client. + """ + # Quick check: Only create a new client if we do not already have one. + if self._operations_client is None: + self._operations_client = operations_v1.OperationsAsyncClient( + self._logged_channel + ) + + # Return the client from cache. + return self._operations_client + @property def evaluate_instances( self, @@ -361,6 +379,34 @@ def evaluate_instances( ) return self._stubs["evaluate_instances"] + @property + def evaluate_dataset( + self, + ) -> Callable[ + [evaluation_service.EvaluateDatasetRequest], Awaitable[operations_pb2.Operation] + ]: + r"""Return a callable for the evaluate dataset method over gRPC. + + Evaluates a dataset based on a set of given metrics. + + Returns: + Callable[[~.EvaluateDatasetRequest], + Awaitable[~.Operation]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "evaluate_dataset" not in self._stubs: + self._stubs["evaluate_dataset"] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1beta1.EvaluationService/EvaluateDataset", + request_serializer=evaluation_service.EvaluateDatasetRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["evaluate_dataset"] + def _prep_wrapped_messages(self, client_info): """Precompute the wrapped methods, overriding the base class method to use async wrappers.""" self._wrapped_methods = { @@ -369,6 +415,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=60.0, client_info=client_info, ), + self.evaluate_dataset: self._wrap_method( + self.evaluate_dataset, + default_timeout=None, + client_info=client_info, + ), self.get_location: self._wrap_method( self.get_location, default_timeout=None, diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest.py index 411f22028a..8416a94dae 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest.py @@ -25,6 +25,7 @@ from google.api_core import gapic_v1 from google.protobuf import json_format +from google.api_core import operations_v1 from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore from google.cloud.location import locations_pb2 # type: ignore @@ -78,6 +79,14 @@ class EvaluationServiceRestInterceptor: .. code-block:: python class MyCustomEvaluationServiceInterceptor(EvaluationServiceRestInterceptor): + def pre_evaluate_dataset(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_evaluate_dataset(self, response): + logging.log(f"Received response: {response}") + return response + def pre_evaluate_instances(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -92,6 +101,55 @@ def post_evaluate_instances(self, response): """ + def pre_evaluate_dataset( + self, + request: evaluation_service.EvaluateDatasetRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateDatasetRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for evaluate_dataset + + Override in a subclass to manipulate the request or metadata + before they are sent to the EvaluationService server. + """ + return request, metadata + + def post_evaluate_dataset( + self, response: operations_pb2.Operation + ) -> operations_pb2.Operation: + """Post-rpc interceptor for evaluate_dataset + + DEPRECATED. Please use the `post_evaluate_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the EvaluationService server but before + it is returned to user code. This `post_evaluate_dataset` interceptor runs + before the `post_evaluate_dataset_with_metadata` interceptor. + """ + return response + + def post_evaluate_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for evaluate_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_dataset_with_metadata` + interceptor in new development instead of the `post_evaluate_dataset` interceptor. + When both interceptors are used, this `post_evaluate_dataset_with_metadata` interceptor runs after the + `post_evaluate_dataset` interceptor. The (possibly modified) response returned by + `post_evaluate_dataset` will be passed to + `post_evaluate_dataset_with_metadata`. + """ + return response, metadata + def pre_evaluate_instances( self, request: evaluation_service.EvaluateInstancesRequest, @@ -112,12 +170,38 @@ def post_evaluate_instances( ) -> evaluation_service.EvaluateInstancesResponse: """Post-rpc interceptor for evaluate_instances - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_evaluate_instances_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EvaluationService server but before - it is returned to user code. + it is returned to user code. This `post_evaluate_instances` interceptor runs + before the `post_evaluate_instances_with_metadata` interceptor. """ return response + def post_evaluate_instances_with_metadata( + self, + response: evaluation_service.EvaluateInstancesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateInstancesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for evaluate_instances + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_instances_with_metadata` + interceptor in new development instead of the `post_evaluate_instances` interceptor. + When both interceptors are used, this `post_evaluate_instances_with_metadata` interceptor runs after the + `post_evaluate_instances` interceptor. The (possibly modified) response returned by + `post_evaluate_instances` will be passed to + `post_evaluate_instances_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -443,165 +527,2422 @@ def __init__( self._session = AuthorizedSession( self._credentials, default_host=self.DEFAULT_HOST ) + self._operations_client: Optional[operations_v1.AbstractOperationsClient] = None if client_cert_source_for_mtls: self._session.configure_mtls_channel(client_cert_source_for_mtls) self._interceptor = interceptor or EvaluationServiceRestInterceptor() self._prep_wrapped_messages(client_info) - class _EvaluateInstances( - _BaseEvaluationServiceRestTransport._BaseEvaluateInstances, - EvaluationServiceRestStub, - ): - def __hash__(self): - return hash("EvaluationServiceRestTransport.EvaluateInstances") - - @staticmethod - def _get_response( - host, - metadata, - query_params, - session, - timeout, - transcoded_request, - body=None, - ): - - uri = transcoded_request["uri"] - method = transcoded_request["method"] - headers = dict(metadata) - headers["Content-Type"] = "application/json" - response = getattr(session, method)( - "{host}{uri}".format(host=host, uri=uri), - timeout=timeout, - headers=headers, - params=rest_helpers.flatten_query_params(query_params, strict=True), - data=body, - ) - return response - - def __call__( - self, - request: evaluation_service.EvaluateInstancesRequest, - *, - retry: OptionalRetry = gapic_v1.method.DEFAULT, - timeout: Optional[float] = None, - metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), - ) -> evaluation_service.EvaluateInstancesResponse: - r"""Call the evaluate instances method over HTTP. - - Args: - request (~.evaluation_service.EvaluateInstancesRequest): - The request object. Request message for - EvaluationService.EvaluateInstances. - retry (google.api_core.retry.Retry): Designation of what errors, if any, - should be retried. - timeout (float): The timeout for this request. - metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be - sent along with the request as metadata. Normally, each value must be of type `str`, - but for metadata keys ending with the suffix `-bin`, the corresponding values must - be of type `bytes`. - - Returns: - ~.evaluation_service.EvaluateInstancesResponse: - Response message for - EvaluationService.EvaluateInstances. - - """ - - http_options = ( - _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_http_options() - ) - - request, metadata = self._interceptor.pre_evaluate_instances( - request, metadata - ) - transcoded_request = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_transcoded_request( - http_options, request - ) - - body = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_request_body_json( - transcoded_request - ) - - # Jsonify the query params - query_params = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_query_params_json( - transcoded_request - ) + @property + def operations_client(self) -> operations_v1.AbstractOperationsClient: + """Create the client designed to process long-running operations. - if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( - logging.DEBUG - ): # pragma: NO COVER - request_url = "{host}{uri}".format( - host=self._host, uri=transcoded_request["uri"] - ) - method = transcoded_request["method"] - try: - request_payload = type(request).to_json(request) - except: - request_payload = None - http_request = { - "payload": request_payload, - "requestMethod": method, - "requestUrl": request_url, - "headers": dict(metadata), - } - _LOGGER.debug( - f"Sending request for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.EvaluateInstances", - extra={ - "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", - "rpcName": "EvaluateInstances", - "httpRequest": http_request, - "metadata": http_request["headers"], + This property caches on the instance; repeated calls return the same + client. + """ + # Only create a new client if we do not already have one. + if self._operations_client is None: + http_options: Dict[str, List[Dict[str, str]]] = { + "google.longrunning.Operations.CancelOperation": [ + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/operations/*}:cancel", }, - ) - - # Send the request - response = EvaluationServiceRestTransport._EvaluateInstances._get_response( - self._host, - metadata, - query_params, - self._session, - timeout, - transcoded_request, - body, - ) - - # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception - # subclass. - if response.status_code >= 400: - raise core_exceptions.from_http_response(response) - - # Return the response - resp = evaluation_service.EvaluateInstancesResponse() - pb_resp = evaluation_service.EvaluateInstancesResponse.pb(resp) - - json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) - - resp = self._interceptor.post_evaluate_instances(resp) - if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( - logging.DEBUG - ): # pragma: NO COVER - try: - response_payload = ( - evaluation_service.EvaluateInstancesResponse.to_json(response) - ) - except: - response_payload = None - http_response = { - "payload": response_payload, - "headers": dict(response.headers), - "status": response.status_code, - } - _LOGGER.debug( - "Received response for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.evaluate_instances", - extra={ - "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", - "rpcName": "EvaluateInstances", - "metadata": http_response["headers"], - "httpResponse": http_response, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}:cancel", }, - ) - return resp + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:cancel", + }, + ], + "google.longrunning.Operations.DeleteOperation": [ + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + ], + "google.longrunning.Operations.GetOperation": [ + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDeploymentJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + ], + "google.longrunning.Operations.ListOperations": [ + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/agents/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/apps/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/schedules/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*}/operations", + }, + ], + "google.longrunning.Operations.WaitOperation": [ + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + ], + } + + rest_transport = operations_v1.OperationsRestTransport( + host=self._host, + # use the credentials which are saved + credentials=self._credentials, + scopes=self._scopes, + http_options=http_options, + path_prefix="v1beta1", + ) + + self._operations_client = operations_v1.AbstractOperationsClient( + transport=rest_transport + ) + + # Return the client from cache. + return self._operations_client + + class _EvaluateDataset( + _BaseEvaluationServiceRestTransport._BaseEvaluateDataset, + EvaluationServiceRestStub, + ): + def __hash__(self): + return hash("EvaluationServiceRestTransport.EvaluateDataset") + + @staticmethod + def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=body, + ) + return response + + def __call__( + self, + request: evaluation_service.EvaluateDatasetRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> operations_pb2.Operation: + r"""Call the evaluate dataset method over HTTP. + + Args: + request (~.evaluation_service.EvaluateDatasetRequest): + The request object. Request message for + EvaluationService.EvaluateDataset. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.operations_pb2.Operation: + This resource represents a + long-running operation that is the + result of a network API call. + + """ + + http_options = ( + _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_http_options() + ) + + request, metadata = self._interceptor.pre_evaluate_dataset( + request, metadata + ) + transcoded_request = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_transcoded_request( + http_options, request + ) + + body = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_request_body_json( + transcoded_request + ) + + # Jsonify the query params + query_params = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = json_format.MessageToJson(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.EvaluateDataset", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateDataset", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = EvaluationServiceRestTransport._EvaluateDataset._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + body, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = operations_pb2.Operation() + json_format.Parse(response.content, resp, ignore_unknown_fields=True) + + resp = self._interceptor.post_evaluate_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_evaluate_dataset_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = json_format.MessageToJson(resp) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": response.status_code, + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.evaluate_dataset", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateDataset", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + return resp + + class _EvaluateInstances( + _BaseEvaluationServiceRestTransport._BaseEvaluateInstances, + EvaluationServiceRestStub, + ): + def __hash__(self): + return hash("EvaluationServiceRestTransport.EvaluateInstances") + + @staticmethod + def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=body, + ) + return response + + def __call__( + self, + request: evaluation_service.EvaluateInstancesRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> evaluation_service.EvaluateInstancesResponse: + r"""Call the evaluate instances method over HTTP. + + Args: + request (~.evaluation_service.EvaluateInstancesRequest): + The request object. Request message for + EvaluationService.EvaluateInstances. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.evaluation_service.EvaluateInstancesResponse: + Response message for + EvaluationService.EvaluateInstances. + + """ + + http_options = ( + _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_http_options() + ) + + request, metadata = self._interceptor.pre_evaluate_instances( + request, metadata + ) + transcoded_request = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_transcoded_request( + http_options, request + ) + + body = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_request_body_json( + transcoded_request + ) + + # Jsonify the query params + query_params = _BaseEvaluationServiceRestTransport._BaseEvaluateInstances._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = type(request).to_json(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.EvaluateInstances", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateInstances", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = EvaluationServiceRestTransport._EvaluateInstances._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + body, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = evaluation_service.EvaluateInstancesResponse() + pb_resp = evaluation_service.EvaluateInstancesResponse.pb(resp) + + json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) + + resp = self._interceptor.post_evaluate_instances(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_evaluate_instances_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = ( + evaluation_service.EvaluateInstancesResponse.to_json(response) + ) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": response.status_code, + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.evaluate_instances", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateInstances", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + return resp + + @property + def evaluate_dataset( + self, + ) -> Callable[ + [evaluation_service.EvaluateDatasetRequest], operations_pb2.Operation + ]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._EvaluateDataset(self._session, self._host, self._interceptor) # type: ignore @property def evaluate_instances( diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_asyncio.py index e5d868ce52..2cef41c8f1 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_asyncio.py @@ -30,6 +30,7 @@ from google.api_core import exceptions as core_exceptions from google.api_core import gapic_v1 +from google.api_core import operations_v1 from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore from google.cloud.location import locations_pb2 # type: ignore @@ -39,6 +40,7 @@ from google.protobuf import json_format +from google.api_core import operations_v1 from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import policy_pb2 # type: ignore from google.cloud.location import locations_pb2 # type: ignore @@ -95,6 +97,14 @@ class AsyncEvaluationServiceRestInterceptor: .. code-block:: python class MyCustomEvaluationServiceInterceptor(EvaluationServiceRestInterceptor): + async def pre_evaluate_dataset(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + async def post_evaluate_dataset(self, response): + logging.log(f"Received response: {response}") + return response + async def pre_evaluate_instances(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -109,6 +119,55 @@ async def post_evaluate_instances(self, response): """ + async def pre_evaluate_dataset( + self, + request: evaluation_service.EvaluateDatasetRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateDatasetRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for evaluate_dataset + + Override in a subclass to manipulate the request or metadata + before they are sent to the EvaluationService server. + """ + return request, metadata + + async def post_evaluate_dataset( + self, response: operations_pb2.Operation + ) -> operations_pb2.Operation: + """Post-rpc interceptor for evaluate_dataset + + DEPRECATED. Please use the `post_evaluate_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the EvaluationService server but before + it is returned to user code. This `post_evaluate_dataset` interceptor runs + before the `post_evaluate_dataset_with_metadata` interceptor. + """ + return response + + async def post_evaluate_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for evaluate_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_dataset_with_metadata` + interceptor in new development instead of the `post_evaluate_dataset` interceptor. + When both interceptors are used, this `post_evaluate_dataset_with_metadata` interceptor runs after the + `post_evaluate_dataset` interceptor. The (possibly modified) response returned by + `post_evaluate_dataset` will be passed to + `post_evaluate_dataset_with_metadata`. + """ + return response, metadata + async def pre_evaluate_instances( self, request: evaluation_service.EvaluateInstancesRequest, @@ -129,12 +188,38 @@ async def post_evaluate_instances( ) -> evaluation_service.EvaluateInstancesResponse: """Post-rpc interceptor for evaluate_instances - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_evaluate_instances_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the EvaluationService server but before - it is returned to user code. + it is returned to user code. This `post_evaluate_instances` interceptor runs + before the `post_evaluate_instances_with_metadata` interceptor. """ return response + async def post_evaluate_instances_with_metadata( + self, + response: evaluation_service.EvaluateInstancesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + evaluation_service.EvaluateInstancesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for evaluate_instances + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the EvaluationService server but before it is returned to user code. + + We recommend only using this `post_evaluate_instances_with_metadata` + interceptor in new development instead of the `post_evaluate_instances` interceptor. + When both interceptors are used, this `post_evaluate_instances_with_metadata` interceptor runs after the + `post_evaluate_instances` interceptor. The (possibly modified) response returned by + `post_evaluate_instances` will be passed to + `post_evaluate_instances_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -443,6 +528,9 @@ def __init__( self._interceptor = interceptor or AsyncEvaluationServiceRestInterceptor() self._wrap_with_kind = True self._prep_wrapped_messages(client_info) + self._operations_client: Optional[ + operations_v1.AsyncOperationsRestClient + ] = None def _prep_wrapped_messages(self, client_info): """Precompute the wrapped methods, overriding the base class method to use async wrappers.""" @@ -452,6 +540,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=60.0, client_info=client_info, ), + self.evaluate_dataset: self._wrap_method( + self.evaluate_dataset, + default_timeout=None, + client_info=client_info, + ), self.get_location: self._wrap_method( self.get_location, default_timeout=None, @@ -509,6 +602,170 @@ def _wrap_method(self, func, *args, **kwargs): kwargs["kind"] = self.kind return gapic_v1.method_async.wrap_method(func, *args, **kwargs) + class _EvaluateDataset( + _BaseEvaluationServiceRestTransport._BaseEvaluateDataset, + AsyncEvaluationServiceRestStub, + ): + def __hash__(self): + return hash("AsyncEvaluationServiceRestTransport.EvaluateDataset") + + @staticmethod + async def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = await getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=body, + ) + return response + + async def __call__( + self, + request: evaluation_service.EvaluateDatasetRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> operations_pb2.Operation: + r"""Call the evaluate dataset method over HTTP. + + Args: + request (~.evaluation_service.EvaluateDatasetRequest): + The request object. Request message for + EvaluationService.EvaluateDataset. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.operations_pb2.Operation: + This resource represents a + long-running operation that is the + result of a network API call. + + """ + + http_options = ( + _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_http_options() + ) + + request, metadata = await self._interceptor.pre_evaluate_dataset( + request, metadata + ) + transcoded_request = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_transcoded_request( + http_options, request + ) + + body = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_request_body_json( + transcoded_request + ) + + # Jsonify the query params + query_params = _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = json_format.MessageToJson(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1beta1.EvaluationServiceClient.EvaluateDataset", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateDataset", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = await AsyncEvaluationServiceRestTransport._EvaluateDataset._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + body, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + content = await response.read() + payload = json.loads(content.decode("utf-8")) + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + raise core_exceptions.format_http_response_error(response, method, request_url, payload) # type: ignore + + # Return the response + resp = operations_pb2.Operation() + pb_resp = resp + content = await response.read() + json_format.Parse(content, pb_resp, ignore_unknown_fields=True) + resp = await self._interceptor.post_evaluate_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_evaluate_dataset_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = json_format.MessageToJson(resp) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": "OK", # need to obtain this properly + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1beta1.EvaluationServiceAsyncClient.evaluate_dataset", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "rpcName": "EvaluateDataset", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + + return resp + class _EvaluateInstances( _BaseEvaluationServiceRestTransport._BaseEvaluateInstances, AsyncEvaluationServiceRestStub, @@ -644,6 +901,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_evaluate_instances(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_evaluate_instances_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -670,6 +931,2099 @@ async def __call__( return resp + @property + def operations_client(self) -> AsyncOperationsRestClient: + """Create the async client designed to process long-running operations. + + This property caches on the instance; repeated calls return the same + client. + """ + # Only create a new client if we do not already have one. + if self._operations_client is None: + http_options: Dict[str, List[Dict[str, str]]] = { + "google.longrunning.Operations.CancelOperation": [ + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:cancel", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:cancel", + }, + ], + "google.longrunning.Operations.DeleteOperation": [ + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "delete", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + ], + "google.longrunning.Operations.GetOperation": [ + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDeploymentJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}", + }, + ], + "google.longrunning.Operations.ListOperations": [ + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/agents/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/apps/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/schedules/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*}/operations", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/solvers/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*}/operations", + }, + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*}/operations", + }, + ], + "google.longrunning.Operations.WaitOperation": [ + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/agents/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/apps/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/edgeDevices/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/endpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensionControllers/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/extensions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/customJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/modelMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/migratableResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/persistentResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/schedules/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/specialistPools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/ui/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/agents/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/apps/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/savedQueries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/annotationSpecs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/datasets/*/dataItems/*/annotations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/deploymentResourcePools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/edgeDevices/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/endpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/evaluationTasks/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/exampleStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensionControllers/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/extensions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featurestores/*/entityTypes/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/customJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/dataLabelingJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/hyperparameterTuningJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/indexEndpoints/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/artifacts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/contexts/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/metadataStores/*/executions/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelDeploymentMonitoringJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/modelMonitors/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/migratableResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*/evaluations/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookExecutionJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimes/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/notebookRuntimeTemplates/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/persistentResources/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/ragCorpora/*/ragFiles/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/reasoningEngines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/studies/*/trials/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/trainingPipelines/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/pipelineJobs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/schedules/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/specialistPools/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/tensorboards/*/experiments/*/runs/*/timeSeries/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureOnlineStores/*/featureViews/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/features/*/operations/*}:wait", + }, + { + "method": "post", + "uri": "/v1beta1/{name=projects/*/locations/*/featureGroups/*/featureMonitors/*/operations/*}:wait", + }, + ], + } + + rest_transport = operations_v1.AsyncOperationsRestTransport( # type: ignore + host=self._host, + # use the credentials which are saved + credentials=self._credentials, # type: ignore + http_options=http_options, + path_prefix="v1beta1", + ) + + self._operations_client = AsyncOperationsRestClient( + transport=rest_transport + ) + + # Return the client from cache. + return self._operations_client + + @property + def evaluate_dataset( + self, + ) -> Callable[ + [evaluation_service.EvaluateDatasetRequest], operations_pb2.Operation + ]: + return self._EvaluateDataset(self._session, self._host, self._interceptor) # type: ignore + @property def evaluate_instances( self, diff --git a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_base.py b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_base.py index 98c8c7aedb..8f2b6414fc 100644 --- a/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_base.py +++ b/google/cloud/aiplatform_v1beta1/services/evaluation_service/transports/rest_base.py @@ -93,6 +93,63 @@ def __init__( api_audience=api_audience, ) + class _BaseEvaluateDataset: + def __hash__(self): # pragma: NO COVER + return NotImplementedError("__hash__ must be implemented.") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {} + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + @staticmethod + def _get_http_options(): + http_options: List[Dict[str, str]] = [ + { + "method": "post", + "uri": "/v1beta1/{location=projects/*/locations/*}:evaluateDataset", + "body": "*", + }, + ] + return http_options + + @staticmethod + def _get_transcoded_request(http_options, request): + pb_request = evaluation_service.EvaluateDatasetRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + return transcoded_request + + @staticmethod + def _get_request_body_json(transcoded_request): + # Jsonify the request body + + body = json_format.MessageToJson( + transcoded_request["body"], use_integers_for_enums=True + ) + return body + + @staticmethod + def _get_query_params_json(transcoded_request): + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update( + _BaseEvaluationServiceRestTransport._BaseEvaluateDataset._get_unset_required_fields( + query_params + ) + ) + + query_params["$alt"] = "json;enum-encoding=int" + return query_params + class _BaseEvaluateInstances: def __hash__(self): # pragma: NO COVER return NotImplementedError("__hash__ must be implemented.") diff --git a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/client.py b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/client.py index 80992c4c26..d17f54ae72 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -527,6 +529,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1056,16 +1085,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1111,16 +1144,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1283,16 +1320,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1404,16 +1445,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1526,16 +1571,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1586,16 +1635,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1641,16 +1694,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1696,16 +1753,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest.py index 6b9c38c56a..a60d11c20e 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest.py @@ -120,12 +120,38 @@ def post_execute_extension( ) -> extension_execution_service.ExecuteExtensionResponse: """Post-rpc interceptor for execute_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_execute_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_execute_extension` interceptor runs + before the `post_execute_extension_with_metadata` interceptor. """ return response + def post_execute_extension_with_metadata( + self, + response: extension_execution_service.ExecuteExtensionResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_execution_service.ExecuteExtensionResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for execute_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionExecutionService server but before it is returned to user code. + + We recommend only using this `post_execute_extension_with_metadata` + interceptor in new development instead of the `post_execute_extension` interceptor. + When both interceptors are used, this `post_execute_extension_with_metadata` interceptor runs after the + `post_execute_extension` interceptor. The (possibly modified) response returned by + `post_execute_extension` will be passed to + `post_execute_extension_with_metadata`. + """ + return response, metadata + def pre_query_extension( self, request: extension_execution_service.QueryExtensionRequest, @@ -146,12 +172,38 @@ def post_query_extension( ) -> extension_execution_service.QueryExtensionResponse: """Post-rpc interceptor for query_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_extension` interceptor runs + before the `post_query_extension_with_metadata` interceptor. """ return response + def post_query_extension_with_metadata( + self, + response: extension_execution_service.QueryExtensionResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_execution_service.QueryExtensionResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_extension_with_metadata` + interceptor in new development instead of the `post_query_extension` interceptor. + When both interceptors are used, this `post_query_extension_with_metadata` interceptor runs after the + `post_query_extension` interceptor. The (possibly modified) response returned by + `post_query_extension` will be passed to + `post_query_extension_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -616,6 +668,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_execute_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_execute_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -773,6 +829,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_query_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest_asyncio.py index dbbf8fc075..95d7bab9ef 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_execution_service/transports/rest_asyncio.py @@ -137,12 +137,38 @@ async def post_execute_extension( ) -> extension_execution_service.ExecuteExtensionResponse: """Post-rpc interceptor for execute_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_execute_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_execute_extension` interceptor runs + before the `post_execute_extension_with_metadata` interceptor. """ return response + async def post_execute_extension_with_metadata( + self, + response: extension_execution_service.ExecuteExtensionResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_execution_service.ExecuteExtensionResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for execute_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionExecutionService server but before it is returned to user code. + + We recommend only using this `post_execute_extension_with_metadata` + interceptor in new development instead of the `post_execute_extension` interceptor. + When both interceptors are used, this `post_execute_extension_with_metadata` interceptor runs after the + `post_execute_extension` interceptor. The (possibly modified) response returned by + `post_execute_extension` will be passed to + `post_execute_extension_with_metadata`. + """ + return response, metadata + async def pre_query_extension( self, request: extension_execution_service.QueryExtensionRequest, @@ -163,12 +189,38 @@ async def post_query_extension( ) -> extension_execution_service.QueryExtensionResponse: """Post-rpc interceptor for query_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_extension` interceptor runs + before the `post_query_extension_with_metadata` interceptor. """ return response + async def post_query_extension_with_metadata( + self, + response: extension_execution_service.QueryExtensionResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_execution_service.QueryExtensionResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_extension_with_metadata` + interceptor in new development instead of the `post_query_extension` interceptor. + When both interceptors are used, this `post_query_extension_with_metadata` interceptor runs after the + `post_query_extension` interceptor. The (possibly modified) response returned by + `post_query_extension` will be passed to + `post_query_extension_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -687,6 +739,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_execute_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_execute_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -850,6 +906,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_query_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/client.py b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/client.py index 3f3056bc6e..b8d661b751 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -558,6 +560,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1481,16 +1510,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1536,16 +1569,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1708,16 +1745,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1829,16 +1870,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1951,16 +1996,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2011,16 +2060,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2066,16 +2119,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2121,16 +2178,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest.py index fba8c2138d..aacf7239e5 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest.py @@ -147,12 +147,35 @@ def post_delete_extension( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_extension` interceptor runs + before the `post_delete_extension_with_metadata` interceptor. """ return response + def post_delete_extension_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_extension_with_metadata` + interceptor in new development instead of the `post_delete_extension` interceptor. + When both interceptors are used, this `post_delete_extension_with_metadata` interceptor runs after the + `post_delete_extension` interceptor. The (possibly modified) response returned by + `post_delete_extension` will be passed to + `post_delete_extension_with_metadata`. + """ + return response, metadata + def pre_get_extension( self, request: extension_registry_service.GetExtensionRequest, @@ -171,12 +194,35 @@ def pre_get_extension( def post_get_extension(self, response: extension.Extension) -> extension.Extension: """Post-rpc interceptor for get_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_extension` interceptor runs + before the `post_get_extension_with_metadata` interceptor. """ return response + def post_get_extension_with_metadata( + self, + response: extension.Extension, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[extension.Extension, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_extension_with_metadata` + interceptor in new development instead of the `post_get_extension` interceptor. + When both interceptors are used, this `post_get_extension_with_metadata` interceptor runs after the + `post_get_extension` interceptor. The (possibly modified) response returned by + `post_get_extension` will be passed to + `post_get_extension_with_metadata`. + """ + return response, metadata + def pre_import_extension( self, request: extension_registry_service.ImportExtensionRequest, @@ -197,12 +243,35 @@ def post_import_extension( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_import_extension` interceptor runs + before the `post_import_extension_with_metadata` interceptor. """ return response + def post_import_extension_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_import_extension_with_metadata` + interceptor in new development instead of the `post_import_extension` interceptor. + When both interceptors are used, this `post_import_extension_with_metadata` interceptor runs after the + `post_import_extension` interceptor. The (possibly modified) response returned by + `post_import_extension` will be passed to + `post_import_extension_with_metadata`. + """ + return response, metadata + def pre_list_extensions( self, request: extension_registry_service.ListExtensionsRequest, @@ -223,12 +292,38 @@ def post_list_extensions( ) -> extension_registry_service.ListExtensionsResponse: """Post-rpc interceptor for list_extensions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_extensions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_extensions` interceptor runs + before the `post_list_extensions_with_metadata` interceptor. """ return response + def post_list_extensions_with_metadata( + self, + response: extension_registry_service.ListExtensionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_registry_service.ListExtensionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_extensions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_extensions_with_metadata` + interceptor in new development instead of the `post_list_extensions` interceptor. + When both interceptors are used, this `post_list_extensions_with_metadata` interceptor runs after the + `post_list_extensions` interceptor. The (possibly modified) response returned by + `post_list_extensions` will be passed to + `post_list_extensions_with_metadata`. + """ + return response, metadata + def pre_update_extension( self, request: extension_registry_service.UpdateExtensionRequest, @@ -249,12 +344,35 @@ def post_update_extension( ) -> gca_extension.Extension: """Post-rpc interceptor for update_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_extension` interceptor runs + before the `post_update_extension_with_metadata` interceptor. """ return response + def post_update_extension_with_metadata( + self, + response: gca_extension.Extension, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_extension.Extension, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_extension_with_metadata` + interceptor in new development instead of the `post_update_extension` interceptor. + When both interceptors are used, this `post_update_extension_with_metadata` interceptor runs after the + `post_update_extension` interceptor. The (possibly modified) response returned by + `post_update_extension` will be passed to + `post_update_extension_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2797,6 +2915,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2945,6 +3067,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3099,6 +3225,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3246,6 +3376,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_extensions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_extensions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3406,6 +3540,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest_asyncio.py index 90bff0f4b9..4936e8bf78 100644 --- a/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/extension_registry_service/transports/rest_asyncio.py @@ -165,12 +165,35 @@ async def post_delete_extension( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_extension` interceptor runs + before the `post_delete_extension_with_metadata` interceptor. """ return response + async def post_delete_extension_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_extension_with_metadata` + interceptor in new development instead of the `post_delete_extension` interceptor. + When both interceptors are used, this `post_delete_extension_with_metadata` interceptor runs after the + `post_delete_extension` interceptor. The (possibly modified) response returned by + `post_delete_extension` will be passed to + `post_delete_extension_with_metadata`. + """ + return response, metadata + async def pre_get_extension( self, request: extension_registry_service.GetExtensionRequest, @@ -191,12 +214,35 @@ async def post_get_extension( ) -> extension.Extension: """Post-rpc interceptor for get_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_extension` interceptor runs + before the `post_get_extension_with_metadata` interceptor. """ return response + async def post_get_extension_with_metadata( + self, + response: extension.Extension, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[extension.Extension, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_extension_with_metadata` + interceptor in new development instead of the `post_get_extension` interceptor. + When both interceptors are used, this `post_get_extension_with_metadata` interceptor runs after the + `post_get_extension` interceptor. The (possibly modified) response returned by + `post_get_extension` will be passed to + `post_get_extension_with_metadata`. + """ + return response, metadata + async def pre_import_extension( self, request: extension_registry_service.ImportExtensionRequest, @@ -217,12 +263,35 @@ async def post_import_extension( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_import_extension` interceptor runs + before the `post_import_extension_with_metadata` interceptor. """ return response + async def post_import_extension_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_import_extension_with_metadata` + interceptor in new development instead of the `post_import_extension` interceptor. + When both interceptors are used, this `post_import_extension_with_metadata` interceptor runs after the + `post_import_extension` interceptor. The (possibly modified) response returned by + `post_import_extension` will be passed to + `post_import_extension_with_metadata`. + """ + return response, metadata + async def pre_list_extensions( self, request: extension_registry_service.ListExtensionsRequest, @@ -243,12 +312,38 @@ async def post_list_extensions( ) -> extension_registry_service.ListExtensionsResponse: """Post-rpc interceptor for list_extensions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_extensions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_extensions` interceptor runs + before the `post_list_extensions_with_metadata` interceptor. """ return response + async def post_list_extensions_with_metadata( + self, + response: extension_registry_service.ListExtensionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + extension_registry_service.ListExtensionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_extensions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_extensions_with_metadata` + interceptor in new development instead of the `post_list_extensions` interceptor. + When both interceptors are used, this `post_list_extensions_with_metadata` interceptor runs after the + `post_list_extensions` interceptor. The (possibly modified) response returned by + `post_list_extensions` will be passed to + `post_list_extensions_with_metadata`. + """ + return response, metadata + async def pre_update_extension( self, request: extension_registry_service.UpdateExtensionRequest, @@ -269,12 +364,35 @@ async def post_update_extension( ) -> gca_extension.Extension: """Post-rpc interceptor for update_extension - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_extension_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ExtensionRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_extension` interceptor runs + before the `post_update_extension_with_metadata` interceptor. """ return response + async def post_update_extension_with_metadata( + self, + response: gca_extension.Extension, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_extension.Extension, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_extension + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ExtensionRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_extension_with_metadata` + interceptor in new development instead of the `post_update_extension` interceptor. + When both interceptors are used, this `post_update_extension_with_metadata` interceptor runs after the + `post_update_extension` interceptor. The (possibly modified) response returned by + `post_update_extension` will be passed to + `post_update_extension_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -806,6 +924,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -960,6 +1082,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1120,6 +1246,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1273,6 +1403,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_extensions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_extensions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1437,6 +1571,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_extension(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_extension_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/client.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/client.py index 7969f6ca42..8318eb9158 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -582,6 +584,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2653,16 +2682,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2708,16 +2741,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2880,16 +2917,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3001,16 +3042,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3123,16 +3168,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3183,16 +3232,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3238,16 +3291,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3293,16 +3350,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest.py index 52f7825304..b9f824471c 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest.py @@ -212,12 +212,35 @@ def post_create_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_online_store` interceptor runs + before the `post_create_feature_online_store_with_metadata` interceptor. """ return response + def post_create_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_online_store_with_metadata` + interceptor in new development instead of the `post_create_feature_online_store` interceptor. + When both interceptors are used, this `post_create_feature_online_store_with_metadata` interceptor runs after the + `post_create_feature_online_store` interceptor. The (possibly modified) response returned by + `post_create_feature_online_store` will be passed to + `post_create_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_create_feature_view( self, request: feature_online_store_admin_service.CreateFeatureViewRequest, @@ -238,12 +261,35 @@ def post_create_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_view` interceptor runs + before the `post_create_feature_view_with_metadata` interceptor. """ return response + def post_create_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_view_with_metadata` + interceptor in new development instead of the `post_create_feature_view` interceptor. + When both interceptors are used, this `post_create_feature_view_with_metadata` interceptor runs after the + `post_create_feature_view` interceptor. The (possibly modified) response returned by + `post_create_feature_view` will be passed to + `post_create_feature_view_with_metadata`. + """ + return response, metadata + def pre_delete_feature_online_store( self, request: feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest, @@ -264,12 +310,35 @@ def post_delete_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_online_store` interceptor runs + before the `post_delete_feature_online_store_with_metadata` interceptor. """ return response + def post_delete_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_online_store_with_metadata` + interceptor in new development instead of the `post_delete_feature_online_store` interceptor. + When both interceptors are used, this `post_delete_feature_online_store_with_metadata` interceptor runs after the + `post_delete_feature_online_store` interceptor. The (possibly modified) response returned by + `post_delete_feature_online_store` will be passed to + `post_delete_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_delete_feature_view( self, request: feature_online_store_admin_service.DeleteFeatureViewRequest, @@ -290,12 +359,35 @@ def post_delete_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_view` interceptor runs + before the `post_delete_feature_view_with_metadata` interceptor. """ return response + def post_delete_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_view_with_metadata` + interceptor in new development instead of the `post_delete_feature_view` interceptor. + When both interceptors are used, this `post_delete_feature_view_with_metadata` interceptor runs after the + `post_delete_feature_view` interceptor. The (possibly modified) response returned by + `post_delete_feature_view` will be passed to + `post_delete_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_feature_online_store( self, request: feature_online_store_admin_service.GetFeatureOnlineStoreRequest, @@ -316,12 +408,37 @@ def post_get_feature_online_store( ) -> feature_online_store.FeatureOnlineStore: """Post-rpc interceptor for get_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_online_store` interceptor runs + before the `post_get_feature_online_store_with_metadata` interceptor. """ return response + def post_get_feature_online_store_with_metadata( + self, + response: feature_online_store.FeatureOnlineStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store.FeatureOnlineStore, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_online_store_with_metadata` + interceptor in new development instead of the `post_get_feature_online_store` interceptor. + When both interceptors are used, this `post_get_feature_online_store_with_metadata` interceptor runs after the + `post_get_feature_online_store` interceptor. The (possibly modified) response returned by + `post_get_feature_online_store` will be passed to + `post_get_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_get_feature_view( self, request: feature_online_store_admin_service.GetFeatureViewRequest, @@ -342,12 +459,35 @@ def post_get_feature_view( ) -> feature_view.FeatureView: """Post-rpc interceptor for get_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view` interceptor runs + before the `post_get_feature_view_with_metadata` interceptor. """ return response + def post_get_feature_view_with_metadata( + self, + response: feature_view.FeatureView, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_view.FeatureView, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_with_metadata` + interceptor in new development instead of the `post_get_feature_view` interceptor. + When both interceptors are used, this `post_get_feature_view_with_metadata` interceptor runs after the + `post_get_feature_view` interceptor. The (possibly modified) response returned by + `post_get_feature_view` will be passed to + `post_get_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_feature_view_sync( self, request: feature_online_store_admin_service.GetFeatureViewSyncRequest, @@ -368,12 +508,37 @@ def post_get_feature_view_sync( ) -> feature_view_sync.FeatureViewSync: """Post-rpc interceptor for get_feature_view_sync - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_sync_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view_sync` interceptor runs + before the `post_get_feature_view_sync_with_metadata` interceptor. """ return response + def post_get_feature_view_sync_with_metadata( + self, + response: feature_view_sync.FeatureViewSync, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_view_sync.FeatureViewSync, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_view_sync + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_sync_with_metadata` + interceptor in new development instead of the `post_get_feature_view_sync` interceptor. + When both interceptors are used, this `post_get_feature_view_sync_with_metadata` interceptor runs after the + `post_get_feature_view_sync` interceptor. The (possibly modified) response returned by + `post_get_feature_view_sync` will be passed to + `post_get_feature_view_sync_with_metadata`. + """ + return response, metadata + def pre_list_feature_online_stores( self, request: feature_online_store_admin_service.ListFeatureOnlineStoresRequest, @@ -395,12 +560,38 @@ def post_list_feature_online_stores( ) -> feature_online_store_admin_service.ListFeatureOnlineStoresResponse: """Post-rpc interceptor for list_feature_online_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_online_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_online_stores` interceptor runs + before the `post_list_feature_online_stores_with_metadata` interceptor. """ return response + def post_list_feature_online_stores_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_online_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_online_stores_with_metadata` + interceptor in new development instead of the `post_list_feature_online_stores` interceptor. + When both interceptors are used, this `post_list_feature_online_stores_with_metadata` interceptor runs after the + `post_list_feature_online_stores` interceptor. The (possibly modified) response returned by + `post_list_feature_online_stores` will be passed to + `post_list_feature_online_stores_with_metadata`. + """ + return response, metadata + def pre_list_feature_views( self, request: feature_online_store_admin_service.ListFeatureViewsRequest, @@ -421,12 +612,38 @@ def post_list_feature_views( ) -> feature_online_store_admin_service.ListFeatureViewsResponse: """Post-rpc interceptor for list_feature_views - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_views_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_views` interceptor runs + before the `post_list_feature_views_with_metadata` interceptor. """ return response + def post_list_feature_views_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_views + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_views_with_metadata` + interceptor in new development instead of the `post_list_feature_views` interceptor. + When both interceptors are used, this `post_list_feature_views_with_metadata` interceptor runs after the + `post_list_feature_views` interceptor. The (possibly modified) response returned by + `post_list_feature_views` will be passed to + `post_list_feature_views_with_metadata`. + """ + return response, metadata + def pre_list_feature_view_syncs( self, request: feature_online_store_admin_service.ListFeatureViewSyncsRequest, @@ -447,12 +664,38 @@ def post_list_feature_view_syncs( ) -> feature_online_store_admin_service.ListFeatureViewSyncsResponse: """Post-rpc interceptor for list_feature_view_syncs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_view_syncs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_view_syncs` interceptor runs + before the `post_list_feature_view_syncs_with_metadata` interceptor. """ return response + def post_list_feature_view_syncs_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewSyncsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewSyncsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_view_syncs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_view_syncs_with_metadata` + interceptor in new development instead of the `post_list_feature_view_syncs` interceptor. + When both interceptors are used, this `post_list_feature_view_syncs_with_metadata` interceptor runs after the + `post_list_feature_view_syncs` interceptor. The (possibly modified) response returned by + `post_list_feature_view_syncs` will be passed to + `post_list_feature_view_syncs_with_metadata`. + """ + return response, metadata + def pre_sync_feature_view( self, request: feature_online_store_admin_service.SyncFeatureViewRequest, @@ -473,12 +716,38 @@ def post_sync_feature_view( ) -> feature_online_store_admin_service.SyncFeatureViewResponse: """Post-rpc interceptor for sync_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_sync_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_sync_feature_view` interceptor runs + before the `post_sync_feature_view_with_metadata` interceptor. """ return response + def post_sync_feature_view_with_metadata( + self, + response: feature_online_store_admin_service.SyncFeatureViewResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.SyncFeatureViewResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for sync_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_sync_feature_view_with_metadata` + interceptor in new development instead of the `post_sync_feature_view` interceptor. + When both interceptors are used, this `post_sync_feature_view_with_metadata` interceptor runs after the + `post_sync_feature_view` interceptor. The (possibly modified) response returned by + `post_sync_feature_view` will be passed to + `post_sync_feature_view_with_metadata`. + """ + return response, metadata + def pre_update_feature_online_store( self, request: feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest, @@ -499,12 +768,35 @@ def post_update_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_online_store` interceptor runs + before the `post_update_feature_online_store_with_metadata` interceptor. """ return response + def post_update_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_online_store_with_metadata` + interceptor in new development instead of the `post_update_feature_online_store` interceptor. + When both interceptors are used, this `post_update_feature_online_store_with_metadata` interceptor runs after the + `post_update_feature_online_store` interceptor. The (possibly modified) response returned by + `post_update_feature_online_store` will be passed to + `post_update_feature_online_store_with_metadata`. + """ + return response, metadata + def pre_update_feature_view( self, request: feature_online_store_admin_service.UpdateFeatureViewRequest, @@ -525,12 +817,35 @@ def post_update_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_view` interceptor runs + before the `post_update_feature_view_with_metadata` interceptor. """ return response + def post_update_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_view_with_metadata` + interceptor in new development instead of the `post_update_feature_view` interceptor. + When both interceptors are used, this `post_update_feature_view_with_metadata` interceptor runs after the + `post_update_feature_view` interceptor. The (possibly modified) response returned by + `post_update_feature_view` will be passed to + `post_update_feature_view_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3085,6 +3400,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3237,6 +3556,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3386,6 +3709,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3532,6 +3859,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3684,6 +4015,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3834,6 +4169,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3985,6 +4324,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_view_sync(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_view_sync_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4141,6 +4484,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_online_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_online_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4292,6 +4639,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_views(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_views_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4445,6 +4796,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_view_syncs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_view_syncs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4602,6 +4957,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_sync_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_sync_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4759,6 +5118,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4911,6 +5274,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest_asyncio.py index 5c95aaccf8..b160c46355 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_admin_service/transports/rest_asyncio.py @@ -230,12 +230,35 @@ async def post_create_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_online_store` interceptor runs + before the `post_create_feature_online_store_with_metadata` interceptor. """ return response + async def post_create_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_online_store_with_metadata` + interceptor in new development instead of the `post_create_feature_online_store` interceptor. + When both interceptors are used, this `post_create_feature_online_store_with_metadata` interceptor runs after the + `post_create_feature_online_store` interceptor. The (possibly modified) response returned by + `post_create_feature_online_store` will be passed to + `post_create_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_create_feature_view( self, request: feature_online_store_admin_service.CreateFeatureViewRequest, @@ -256,12 +279,35 @@ async def post_create_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_view` interceptor runs + before the `post_create_feature_view_with_metadata` interceptor. """ return response + async def post_create_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_create_feature_view_with_metadata` + interceptor in new development instead of the `post_create_feature_view` interceptor. + When both interceptors are used, this `post_create_feature_view_with_metadata` interceptor runs after the + `post_create_feature_view` interceptor. The (possibly modified) response returned by + `post_create_feature_view` will be passed to + `post_create_feature_view_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_online_store( self, request: feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest, @@ -282,12 +328,35 @@ async def post_delete_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_online_store` interceptor runs + before the `post_delete_feature_online_store_with_metadata` interceptor. """ return response + async def post_delete_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_online_store_with_metadata` + interceptor in new development instead of the `post_delete_feature_online_store` interceptor. + When both interceptors are used, this `post_delete_feature_online_store_with_metadata` interceptor runs after the + `post_delete_feature_online_store` interceptor. The (possibly modified) response returned by + `post_delete_feature_online_store` will be passed to + `post_delete_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_view( self, request: feature_online_store_admin_service.DeleteFeatureViewRequest, @@ -308,12 +377,35 @@ async def post_delete_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_view` interceptor runs + before the `post_delete_feature_view_with_metadata` interceptor. """ return response + async def post_delete_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_view_with_metadata` + interceptor in new development instead of the `post_delete_feature_view` interceptor. + When both interceptors are used, this `post_delete_feature_view_with_metadata` interceptor runs after the + `post_delete_feature_view` interceptor. The (possibly modified) response returned by + `post_delete_feature_view` will be passed to + `post_delete_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_feature_online_store( self, request: feature_online_store_admin_service.GetFeatureOnlineStoreRequest, @@ -334,12 +426,37 @@ async def post_get_feature_online_store( ) -> feature_online_store.FeatureOnlineStore: """Post-rpc interceptor for get_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_online_store` interceptor runs + before the `post_get_feature_online_store_with_metadata` interceptor. """ return response + async def post_get_feature_online_store_with_metadata( + self, + response: feature_online_store.FeatureOnlineStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store.FeatureOnlineStore, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_online_store_with_metadata` + interceptor in new development instead of the `post_get_feature_online_store` interceptor. + When both interceptors are used, this `post_get_feature_online_store_with_metadata` interceptor runs after the + `post_get_feature_online_store` interceptor. The (possibly modified) response returned by + `post_get_feature_online_store` will be passed to + `post_get_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_get_feature_view( self, request: feature_online_store_admin_service.GetFeatureViewRequest, @@ -360,12 +477,35 @@ async def post_get_feature_view( ) -> feature_view.FeatureView: """Post-rpc interceptor for get_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view` interceptor runs + before the `post_get_feature_view_with_metadata` interceptor. """ return response + async def post_get_feature_view_with_metadata( + self, + response: feature_view.FeatureView, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_view.FeatureView, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_with_metadata` + interceptor in new development instead of the `post_get_feature_view` interceptor. + When both interceptors are used, this `post_get_feature_view_with_metadata` interceptor runs after the + `post_get_feature_view` interceptor. The (possibly modified) response returned by + `post_get_feature_view` will be passed to + `post_get_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_feature_view_sync( self, request: feature_online_store_admin_service.GetFeatureViewSyncRequest, @@ -386,12 +526,37 @@ async def post_get_feature_view_sync( ) -> feature_view_sync.FeatureViewSync: """Post-rpc interceptor for get_feature_view_sync - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_view_sync_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_view_sync` interceptor runs + before the `post_get_feature_view_sync_with_metadata` interceptor. """ return response + async def post_get_feature_view_sync_with_metadata( + self, + response: feature_view_sync.FeatureViewSync, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_view_sync.FeatureViewSync, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_view_sync + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_get_feature_view_sync_with_metadata` + interceptor in new development instead of the `post_get_feature_view_sync` interceptor. + When both interceptors are used, this `post_get_feature_view_sync_with_metadata` interceptor runs after the + `post_get_feature_view_sync` interceptor. The (possibly modified) response returned by + `post_get_feature_view_sync` will be passed to + `post_get_feature_view_sync_with_metadata`. + """ + return response, metadata + async def pre_list_feature_online_stores( self, request: feature_online_store_admin_service.ListFeatureOnlineStoresRequest, @@ -413,12 +578,38 @@ async def post_list_feature_online_stores( ) -> feature_online_store_admin_service.ListFeatureOnlineStoresResponse: """Post-rpc interceptor for list_feature_online_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_online_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_online_stores` interceptor runs + before the `post_list_feature_online_stores_with_metadata` interceptor. """ return response + async def post_list_feature_online_stores_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureOnlineStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_online_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_online_stores_with_metadata` + interceptor in new development instead of the `post_list_feature_online_stores` interceptor. + When both interceptors are used, this `post_list_feature_online_stores_with_metadata` interceptor runs after the + `post_list_feature_online_stores` interceptor. The (possibly modified) response returned by + `post_list_feature_online_stores` will be passed to + `post_list_feature_online_stores_with_metadata`. + """ + return response, metadata + async def pre_list_feature_views( self, request: feature_online_store_admin_service.ListFeatureViewsRequest, @@ -439,12 +630,38 @@ async def post_list_feature_views( ) -> feature_online_store_admin_service.ListFeatureViewsResponse: """Post-rpc interceptor for list_feature_views - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_views_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_views` interceptor runs + before the `post_list_feature_views_with_metadata` interceptor. """ return response + async def post_list_feature_views_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_views + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_views_with_metadata` + interceptor in new development instead of the `post_list_feature_views` interceptor. + When both interceptors are used, this `post_list_feature_views_with_metadata` interceptor runs after the + `post_list_feature_views` interceptor. The (possibly modified) response returned by + `post_list_feature_views` will be passed to + `post_list_feature_views_with_metadata`. + """ + return response, metadata + async def pre_list_feature_view_syncs( self, request: feature_online_store_admin_service.ListFeatureViewSyncsRequest, @@ -465,12 +682,38 @@ async def post_list_feature_view_syncs( ) -> feature_online_store_admin_service.ListFeatureViewSyncsResponse: """Post-rpc interceptor for list_feature_view_syncs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_view_syncs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_view_syncs` interceptor runs + before the `post_list_feature_view_syncs_with_metadata` interceptor. """ return response + async def post_list_feature_view_syncs_with_metadata( + self, + response: feature_online_store_admin_service.ListFeatureViewSyncsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.ListFeatureViewSyncsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_view_syncs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_list_feature_view_syncs_with_metadata` + interceptor in new development instead of the `post_list_feature_view_syncs` interceptor. + When both interceptors are used, this `post_list_feature_view_syncs_with_metadata` interceptor runs after the + `post_list_feature_view_syncs` interceptor. The (possibly modified) response returned by + `post_list_feature_view_syncs` will be passed to + `post_list_feature_view_syncs_with_metadata`. + """ + return response, metadata + async def pre_sync_feature_view( self, request: feature_online_store_admin_service.SyncFeatureViewRequest, @@ -491,12 +734,38 @@ async def post_sync_feature_view( ) -> feature_online_store_admin_service.SyncFeatureViewResponse: """Post-rpc interceptor for sync_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_sync_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_sync_feature_view` interceptor runs + before the `post_sync_feature_view_with_metadata` interceptor. """ return response + async def post_sync_feature_view_with_metadata( + self, + response: feature_online_store_admin_service.SyncFeatureViewResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_admin_service.SyncFeatureViewResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for sync_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_sync_feature_view_with_metadata` + interceptor in new development instead of the `post_sync_feature_view` interceptor. + When both interceptors are used, this `post_sync_feature_view_with_metadata` interceptor runs after the + `post_sync_feature_view` interceptor. The (possibly modified) response returned by + `post_sync_feature_view` will be passed to + `post_sync_feature_view_with_metadata`. + """ + return response, metadata + async def pre_update_feature_online_store( self, request: feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest, @@ -517,12 +786,35 @@ async def post_update_feature_online_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_online_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_online_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_online_store` interceptor runs + before the `post_update_feature_online_store_with_metadata` interceptor. """ return response + async def post_update_feature_online_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_online_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_online_store_with_metadata` + interceptor in new development instead of the `post_update_feature_online_store` interceptor. + When both interceptors are used, this `post_update_feature_online_store_with_metadata` interceptor runs after the + `post_update_feature_online_store` interceptor. The (possibly modified) response returned by + `post_update_feature_online_store` will be passed to + `post_update_feature_online_store_with_metadata`. + """ + return response, metadata + async def pre_update_feature_view( self, request: feature_online_store_admin_service.UpdateFeatureViewRequest, @@ -543,12 +835,35 @@ async def post_update_feature_view( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_view - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_view_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreAdminService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_view` interceptor runs + before the `post_update_feature_view_with_metadata` interceptor. """ return response + async def post_update_feature_view_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_view + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreAdminService server but before it is returned to user code. + + We recommend only using this `post_update_feature_view_with_metadata` + interceptor in new development instead of the `post_update_feature_view` interceptor. + When both interceptors are used, this `post_update_feature_view_with_metadata` interceptor runs after the + `post_update_feature_view` interceptor. The (possibly modified) response returned by + `post_update_feature_view` will be passed to + `post_update_feature_view_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1132,6 +1447,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1294,6 +1616,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1451,6 +1777,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1607,6 +1940,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1765,6 +2102,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1923,6 +2267,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2080,6 +2428,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_view_sync(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_view_sync_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2242,6 +2594,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_online_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_feature_online_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2401,6 +2760,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_views(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_feature_views_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2560,6 +2923,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_view_syncs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_feature_view_syncs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2725,6 +3095,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_sync_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_sync_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2890,6 +3264,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_online_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_feature_online_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3052,6 +3433,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_view(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_view_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/client.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/client.py index 663c42eadf..21efff178d 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -510,6 +512,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1113,16 +1142,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1168,16 +1201,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1340,16 +1377,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1461,16 +1502,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1583,16 +1628,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1643,16 +1692,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1698,16 +1751,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1753,16 +1810,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest.py index c6fa000bbb..772b5fabcb 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest.py @@ -120,12 +120,38 @@ def post_fetch_feature_values( ) -> feature_online_store_service.FetchFeatureValuesResponse: """Post-rpc interceptor for fetch_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_fetch_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_fetch_feature_values` interceptor runs + before the `post_fetch_feature_values_with_metadata` interceptor. """ return response + def post_fetch_feature_values_with_metadata( + self, + response: feature_online_store_service.FetchFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.FetchFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for fetch_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_fetch_feature_values_with_metadata` + interceptor in new development instead of the `post_fetch_feature_values` interceptor. + When both interceptors are used, this `post_fetch_feature_values_with_metadata` interceptor runs after the + `post_fetch_feature_values` interceptor. The (possibly modified) response returned by + `post_fetch_feature_values` will be passed to + `post_fetch_feature_values_with_metadata`. + """ + return response, metadata + def pre_search_nearest_entities( self, request: feature_online_store_service.SearchNearestEntitiesRequest, @@ -146,12 +172,38 @@ def post_search_nearest_entities( ) -> feature_online_store_service.SearchNearestEntitiesResponse: """Post-rpc interceptor for search_nearest_entities - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_nearest_entities_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_nearest_entities` interceptor runs + before the `post_search_nearest_entities_with_metadata` interceptor. """ return response + def post_search_nearest_entities_with_metadata( + self, + response: feature_online_store_service.SearchNearestEntitiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.SearchNearestEntitiesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_nearest_entities + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_search_nearest_entities_with_metadata` + interceptor in new development instead of the `post_search_nearest_entities` interceptor. + When both interceptors are used, this `post_search_nearest_entities_with_metadata` interceptor runs after the + `post_search_nearest_entities` interceptor. The (possibly modified) response returned by + `post_search_nearest_entities` will be passed to + `post_search_nearest_entities_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -616,6 +668,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_fetch_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_fetch_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -775,6 +831,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_nearest_entities(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_nearest_entities_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest_asyncio.py index 09c1dd4919..2358c73dde 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_online_store_service/transports/rest_asyncio.py @@ -137,12 +137,38 @@ async def post_fetch_feature_values( ) -> feature_online_store_service.FetchFeatureValuesResponse: """Post-rpc interceptor for fetch_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_fetch_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_fetch_feature_values` interceptor runs + before the `post_fetch_feature_values_with_metadata` interceptor. """ return response + async def post_fetch_feature_values_with_metadata( + self, + response: feature_online_store_service.FetchFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.FetchFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for fetch_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_fetch_feature_values_with_metadata` + interceptor in new development instead of the `post_fetch_feature_values` interceptor. + When both interceptors are used, this `post_fetch_feature_values_with_metadata` interceptor runs after the + `post_fetch_feature_values` interceptor. The (possibly modified) response returned by + `post_fetch_feature_values` will be passed to + `post_fetch_feature_values_with_metadata`. + """ + return response, metadata + async def pre_search_nearest_entities( self, request: feature_online_store_service.SearchNearestEntitiesRequest, @@ -163,12 +189,38 @@ async def post_search_nearest_entities( ) -> feature_online_store_service.SearchNearestEntitiesResponse: """Post-rpc interceptor for search_nearest_entities - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_nearest_entities_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureOnlineStoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_nearest_entities` interceptor runs + before the `post_search_nearest_entities_with_metadata` interceptor. """ return response + async def post_search_nearest_entities_with_metadata( + self, + response: feature_online_store_service.SearchNearestEntitiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_online_store_service.SearchNearestEntitiesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_nearest_entities + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureOnlineStoreService server but before it is returned to user code. + + We recommend only using this `post_search_nearest_entities_with_metadata` + interceptor in new development instead of the `post_search_nearest_entities` interceptor. + When both interceptors are used, this `post_search_nearest_entities_with_metadata` interceptor runs after the + `post_search_nearest_entities` interceptor. The (possibly modified) response returned by + `post_search_nearest_entities` will be passed to + `post_search_nearest_entities_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -696,6 +748,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_fetch_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_fetch_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -863,6 +919,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_nearest_entities(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_nearest_entities_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/client.py b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/client.py index 61cee0fe4b..4cfe6b009a 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -601,6 +603,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3409,16 +3438,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3464,16 +3497,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3636,16 +3673,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3757,16 +3798,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3879,16 +3924,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3939,16 +3988,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3994,16 +4047,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -4049,16 +4106,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest.py index ef4bb118dd..62d4dbdb0f 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest.py @@ -265,12 +265,35 @@ def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -291,12 +314,35 @@ def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + def pre_create_feature_group( self, request: feature_registry_service.CreateFeatureGroupRequest, @@ -317,12 +363,35 @@ def post_create_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_group` interceptor runs + before the `post_create_feature_group_with_metadata` interceptor. """ return response + def post_create_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_group_with_metadata` + interceptor in new development instead of the `post_create_feature_group` interceptor. + When both interceptors are used, this `post_create_feature_group_with_metadata` interceptor runs after the + `post_create_feature_group` interceptor. The (possibly modified) response returned by + `post_create_feature_group` will be passed to + `post_create_feature_group_with_metadata`. + """ + return response, metadata + def pre_create_feature_monitor( self, request: feature_registry_service.CreateFeatureMonitorRequest, @@ -343,12 +412,35 @@ def post_create_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_monitor` interceptor runs + before the `post_create_feature_monitor_with_metadata` interceptor. """ return response + def post_create_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_monitor_with_metadata` + interceptor in new development instead of the `post_create_feature_monitor` interceptor. + When both interceptors are used, this `post_create_feature_monitor_with_metadata` interceptor runs after the + `post_create_feature_monitor` interceptor. The (possibly modified) response returned by + `post_create_feature_monitor` will be passed to + `post_create_feature_monitor_with_metadata`. + """ + return response, metadata + def pre_create_feature_monitor_job( self, request: feature_registry_service.CreateFeatureMonitorJobRequest, @@ -369,12 +461,38 @@ def post_create_feature_monitor_job( ) -> gca_feature_monitor_job.FeatureMonitorJob: """Post-rpc interceptor for create_feature_monitor_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_monitor_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_monitor_job` interceptor runs + before the `post_create_feature_monitor_job_with_metadata` interceptor. """ return response + def post_create_feature_monitor_job_with_metadata( + self, + response: gca_feature_monitor_job.FeatureMonitorJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_feature_monitor_job.FeatureMonitorJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_feature_monitor_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_monitor_job_with_metadata` + interceptor in new development instead of the `post_create_feature_monitor_job` interceptor. + When both interceptors are used, this `post_create_feature_monitor_job_with_metadata` interceptor runs after the + `post_create_feature_monitor_job` interceptor. The (possibly modified) response returned by + `post_create_feature_monitor_job` will be passed to + `post_create_feature_monitor_job_with_metadata`. + """ + return response, metadata + def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -395,12 +513,35 @@ def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + def pre_delete_feature_group( self, request: feature_registry_service.DeleteFeatureGroupRequest, @@ -421,12 +562,35 @@ def post_delete_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_group` interceptor runs + before the `post_delete_feature_group_with_metadata` interceptor. """ return response + def post_delete_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_group_with_metadata` + interceptor in new development instead of the `post_delete_feature_group` interceptor. + When both interceptors are used, this `post_delete_feature_group_with_metadata` interceptor runs after the + `post_delete_feature_group` interceptor. The (possibly modified) response returned by + `post_delete_feature_group` will be passed to + `post_delete_feature_group_with_metadata`. + """ + return response, metadata + def pre_delete_feature_monitor( self, request: feature_registry_service.DeleteFeatureMonitorRequest, @@ -447,12 +611,35 @@ def post_delete_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_monitor` interceptor runs + before the `post_delete_feature_monitor_with_metadata` interceptor. """ return response + def post_delete_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_monitor_with_metadata` + interceptor in new development instead of the `post_delete_feature_monitor` interceptor. + When both interceptors are used, this `post_delete_feature_monitor_with_metadata` interceptor runs after the + `post_delete_feature_monitor` interceptor. The (possibly modified) response returned by + `post_delete_feature_monitor` will be passed to + `post_delete_feature_monitor_with_metadata`. + """ + return response, metadata + def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -470,12 +657,35 @@ def pre_get_feature( def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + def pre_get_feature_group( self, request: feature_registry_service.GetFeatureGroupRequest, @@ -496,12 +706,35 @@ def post_get_feature_group( ) -> feature_group.FeatureGroup: """Post-rpc interceptor for get_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_group` interceptor runs + before the `post_get_feature_group_with_metadata` interceptor. """ return response + def post_get_feature_group_with_metadata( + self, + response: feature_group.FeatureGroup, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_group.FeatureGroup, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_group_with_metadata` + interceptor in new development instead of the `post_get_feature_group` interceptor. + When both interceptors are used, this `post_get_feature_group_with_metadata` interceptor runs after the + `post_get_feature_group` interceptor. The (possibly modified) response returned by + `post_get_feature_group` will be passed to + `post_get_feature_group_with_metadata`. + """ + return response, metadata + def pre_get_feature_monitor( self, request: feature_registry_service.GetFeatureMonitorRequest, @@ -522,12 +755,35 @@ def post_get_feature_monitor( ) -> feature_monitor.FeatureMonitor: """Post-rpc interceptor for get_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_monitor` interceptor runs + before the `post_get_feature_monitor_with_metadata` interceptor. """ return response + def post_get_feature_monitor_with_metadata( + self, + response: feature_monitor.FeatureMonitor, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_monitor.FeatureMonitor, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_monitor_with_metadata` + interceptor in new development instead of the `post_get_feature_monitor` interceptor. + When both interceptors are used, this `post_get_feature_monitor_with_metadata` interceptor runs after the + `post_get_feature_monitor` interceptor. The (possibly modified) response returned by + `post_get_feature_monitor` will be passed to + `post_get_feature_monitor_with_metadata`. + """ + return response, metadata + def pre_get_feature_monitor_job( self, request: feature_registry_service.GetFeatureMonitorJobRequest, @@ -548,12 +804,37 @@ def post_get_feature_monitor_job( ) -> feature_monitor_job.FeatureMonitorJob: """Post-rpc interceptor for get_feature_monitor_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_monitor_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_monitor_job` interceptor runs + before the `post_get_feature_monitor_job_with_metadata` interceptor. """ return response + def post_get_feature_monitor_job_with_metadata( + self, + response: feature_monitor_job.FeatureMonitorJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_monitor_job.FeatureMonitorJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_monitor_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_monitor_job_with_metadata` + interceptor in new development instead of the `post_get_feature_monitor_job` interceptor. + When both interceptors are used, this `post_get_feature_monitor_job_with_metadata` interceptor runs after the + `post_get_feature_monitor_job` interceptor. The (possibly modified) response returned by + `post_get_feature_monitor_job` will be passed to + `post_get_feature_monitor_job_with_metadata`. + """ + return response, metadata + def pre_list_feature_groups( self, request: feature_registry_service.ListFeatureGroupsRequest, @@ -574,12 +855,38 @@ def post_list_feature_groups( ) -> feature_registry_service.ListFeatureGroupsResponse: """Post-rpc interceptor for list_feature_groups - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_groups_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_groups` interceptor runs + before the `post_list_feature_groups_with_metadata` interceptor. """ return response + def post_list_feature_groups_with_metadata( + self, + response: feature_registry_service.ListFeatureGroupsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureGroupsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_groups + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_groups_with_metadata` + interceptor in new development instead of the `post_list_feature_groups` interceptor. + When both interceptors are used, this `post_list_feature_groups_with_metadata` interceptor runs after the + `post_list_feature_groups` interceptor. The (possibly modified) response returned by + `post_list_feature_groups` will be passed to + `post_list_feature_groups_with_metadata`. + """ + return response, metadata + def pre_list_feature_monitor_jobs( self, request: feature_registry_service.ListFeatureMonitorJobsRequest, @@ -600,12 +907,38 @@ def post_list_feature_monitor_jobs( ) -> feature_registry_service.ListFeatureMonitorJobsResponse: """Post-rpc interceptor for list_feature_monitor_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_monitor_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_monitor_jobs` interceptor runs + before the `post_list_feature_monitor_jobs_with_metadata` interceptor. """ return response + def post_list_feature_monitor_jobs_with_metadata( + self, + response: feature_registry_service.ListFeatureMonitorJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureMonitorJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_monitor_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_monitor_jobs_with_metadata` + interceptor in new development instead of the `post_list_feature_monitor_jobs` interceptor. + When both interceptors are used, this `post_list_feature_monitor_jobs_with_metadata` interceptor runs after the + `post_list_feature_monitor_jobs` interceptor. The (possibly modified) response returned by + `post_list_feature_monitor_jobs` will be passed to + `post_list_feature_monitor_jobs_with_metadata`. + """ + return response, metadata + def pre_list_feature_monitors( self, request: feature_registry_service.ListFeatureMonitorsRequest, @@ -626,12 +959,38 @@ def post_list_feature_monitors( ) -> feature_registry_service.ListFeatureMonitorsResponse: """Post-rpc interceptor for list_feature_monitors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_monitors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_monitors` interceptor runs + before the `post_list_feature_monitors_with_metadata` interceptor. """ return response + def post_list_feature_monitors_with_metadata( + self, + response: feature_registry_service.ListFeatureMonitorsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureMonitorsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_monitors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_monitors_with_metadata` + interceptor in new development instead of the `post_list_feature_monitors` interceptor. + When both interceptors are used, this `post_list_feature_monitors_with_metadata` interceptor runs after the + `post_list_feature_monitors` interceptor. The (possibly modified) response returned by + `post_list_feature_monitors` will be passed to + `post_list_feature_monitors_with_metadata`. + """ + return response, metadata + def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -652,12 +1011,38 @@ def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -678,12 +1063,35 @@ def post_update_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + def post_update_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + def pre_update_feature_group( self, request: feature_registry_service.UpdateFeatureGroupRequest, @@ -704,12 +1112,35 @@ def post_update_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_group` interceptor runs + before the `post_update_feature_group_with_metadata` interceptor. """ return response + def post_update_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_group_with_metadata` + interceptor in new development instead of the `post_update_feature_group` interceptor. + When both interceptors are used, this `post_update_feature_group_with_metadata` interceptor runs after the + `post_update_feature_group` interceptor. The (possibly modified) response returned by + `post_update_feature_group` will be passed to + `post_update_feature_group_with_metadata`. + """ + return response, metadata + def pre_update_feature_monitor( self, request: feature_registry_service.UpdateFeatureMonitorRequest, @@ -730,12 +1161,35 @@ def post_update_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_monitor` interceptor runs + before the `post_update_feature_monitor_with_metadata` interceptor. """ return response + def post_update_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_monitor_with_metadata` + interceptor in new development instead of the `post_update_feature_monitor` interceptor. + When both interceptors are used, this `post_update_feature_monitor_with_metadata` interceptor runs after the + `post_update_feature_monitor` interceptor. The (possibly modified) response returned by + `post_update_feature_monitor` will be passed to + `post_update_feature_monitor_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3287,6 +3741,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3439,6 +3897,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3593,6 +4055,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3747,6 +4213,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3899,6 +4369,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature_monitor_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_monitor_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4047,6 +4521,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4195,6 +4673,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4343,6 +4825,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4491,6 +4977,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4638,6 +5128,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4785,6 +5279,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4932,6 +5430,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature_monitor_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_monitor_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5083,6 +5585,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_groups(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_groups_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5234,6 +5740,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_monitor_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_monitor_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5387,6 +5897,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_feature_monitors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_feature_monitors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5540,6 +6054,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5694,6 +6212,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5848,6 +6370,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6002,6 +6528,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest_asyncio.py index 9fc411d02e..aef95199cc 100644 --- a/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/feature_registry_service/transports/rest_asyncio.py @@ -283,12 +283,35 @@ async def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + async def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + async def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -309,12 +332,35 @@ async def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + async def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + async def pre_create_feature_group( self, request: feature_registry_service.CreateFeatureGroupRequest, @@ -335,12 +381,35 @@ async def post_create_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_group` interceptor runs + before the `post_create_feature_group_with_metadata` interceptor. """ return response + async def post_create_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_group_with_metadata` + interceptor in new development instead of the `post_create_feature_group` interceptor. + When both interceptors are used, this `post_create_feature_group_with_metadata` interceptor runs after the + `post_create_feature_group` interceptor. The (possibly modified) response returned by + `post_create_feature_group` will be passed to + `post_create_feature_group_with_metadata`. + """ + return response, metadata + async def pre_create_feature_monitor( self, request: feature_registry_service.CreateFeatureMonitorRequest, @@ -361,12 +430,35 @@ async def post_create_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_monitor` interceptor runs + before the `post_create_feature_monitor_with_metadata` interceptor. """ return response + async def post_create_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_monitor_with_metadata` + interceptor in new development instead of the `post_create_feature_monitor` interceptor. + When both interceptors are used, this `post_create_feature_monitor_with_metadata` interceptor runs after the + `post_create_feature_monitor` interceptor. The (possibly modified) response returned by + `post_create_feature_monitor` will be passed to + `post_create_feature_monitor_with_metadata`. + """ + return response, metadata + async def pre_create_feature_monitor_job( self, request: feature_registry_service.CreateFeatureMonitorJobRequest, @@ -387,12 +479,38 @@ async def post_create_feature_monitor_job( ) -> gca_feature_monitor_job.FeatureMonitorJob: """Post-rpc interceptor for create_feature_monitor_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_monitor_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature_monitor_job` interceptor runs + before the `post_create_feature_monitor_job_with_metadata` interceptor. """ return response + async def post_create_feature_monitor_job_with_metadata( + self, + response: gca_feature_monitor_job.FeatureMonitorJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_feature_monitor_job.FeatureMonitorJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_feature_monitor_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_create_feature_monitor_job_with_metadata` + interceptor in new development instead of the `post_create_feature_monitor_job` interceptor. + When both interceptors are used, this `post_create_feature_monitor_job_with_metadata` interceptor runs after the + `post_create_feature_monitor_job` interceptor. The (possibly modified) response returned by + `post_create_feature_monitor_job` will be passed to + `post_create_feature_monitor_job_with_metadata`. + """ + return response, metadata + async def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -413,12 +531,35 @@ async def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + async def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_group( self, request: feature_registry_service.DeleteFeatureGroupRequest, @@ -439,12 +580,35 @@ async def post_delete_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_group` interceptor runs + before the `post_delete_feature_group_with_metadata` interceptor. """ return response + async def post_delete_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_group_with_metadata` + interceptor in new development instead of the `post_delete_feature_group` interceptor. + When both interceptors are used, this `post_delete_feature_group_with_metadata` interceptor runs after the + `post_delete_feature_group` interceptor. The (possibly modified) response returned by + `post_delete_feature_group` will be passed to + `post_delete_feature_group_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_monitor( self, request: feature_registry_service.DeleteFeatureMonitorRequest, @@ -465,12 +629,35 @@ async def post_delete_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_monitor` interceptor runs + before the `post_delete_feature_monitor_with_metadata` interceptor. """ return response + async def post_delete_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_monitor_with_metadata` + interceptor in new development instead of the `post_delete_feature_monitor` interceptor. + When both interceptors are used, this `post_delete_feature_monitor_with_metadata` interceptor runs after the + `post_delete_feature_monitor` interceptor. The (possibly modified) response returned by + `post_delete_feature_monitor` will be passed to + `post_delete_feature_monitor_with_metadata`. + """ + return response, metadata + async def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -488,12 +675,35 @@ async def pre_get_feature( async def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + async def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + async def pre_get_feature_group( self, request: feature_registry_service.GetFeatureGroupRequest, @@ -514,12 +724,35 @@ async def post_get_feature_group( ) -> feature_group.FeatureGroup: """Post-rpc interceptor for get_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_group` interceptor runs + before the `post_get_feature_group_with_metadata` interceptor. """ return response + async def post_get_feature_group_with_metadata( + self, + response: feature_group.FeatureGroup, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_group.FeatureGroup, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_group_with_metadata` + interceptor in new development instead of the `post_get_feature_group` interceptor. + When both interceptors are used, this `post_get_feature_group_with_metadata` interceptor runs after the + `post_get_feature_group` interceptor. The (possibly modified) response returned by + `post_get_feature_group` will be passed to + `post_get_feature_group_with_metadata`. + """ + return response, metadata + async def pre_get_feature_monitor( self, request: feature_registry_service.GetFeatureMonitorRequest, @@ -540,12 +773,35 @@ async def post_get_feature_monitor( ) -> feature_monitor.FeatureMonitor: """Post-rpc interceptor for get_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_monitor` interceptor runs + before the `post_get_feature_monitor_with_metadata` interceptor. """ return response + async def post_get_feature_monitor_with_metadata( + self, + response: feature_monitor.FeatureMonitor, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature_monitor.FeatureMonitor, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_monitor_with_metadata` + interceptor in new development instead of the `post_get_feature_monitor` interceptor. + When both interceptors are used, this `post_get_feature_monitor_with_metadata` interceptor runs after the + `post_get_feature_monitor` interceptor. The (possibly modified) response returned by + `post_get_feature_monitor` will be passed to + `post_get_feature_monitor_with_metadata`. + """ + return response, metadata + async def pre_get_feature_monitor_job( self, request: feature_registry_service.GetFeatureMonitorJobRequest, @@ -566,12 +822,37 @@ async def post_get_feature_monitor_job( ) -> feature_monitor_job.FeatureMonitorJob: """Post-rpc interceptor for get_feature_monitor_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_monitor_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature_monitor_job` interceptor runs + before the `post_get_feature_monitor_job_with_metadata` interceptor. """ return response + async def post_get_feature_monitor_job_with_metadata( + self, + response: feature_monitor_job.FeatureMonitorJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_monitor_job.FeatureMonitorJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_feature_monitor_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_get_feature_monitor_job_with_metadata` + interceptor in new development instead of the `post_get_feature_monitor_job` interceptor. + When both interceptors are used, this `post_get_feature_monitor_job_with_metadata` interceptor runs after the + `post_get_feature_monitor_job` interceptor. The (possibly modified) response returned by + `post_get_feature_monitor_job` will be passed to + `post_get_feature_monitor_job_with_metadata`. + """ + return response, metadata + async def pre_list_feature_groups( self, request: feature_registry_service.ListFeatureGroupsRequest, @@ -592,12 +873,38 @@ async def post_list_feature_groups( ) -> feature_registry_service.ListFeatureGroupsResponse: """Post-rpc interceptor for list_feature_groups - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_groups_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_groups` interceptor runs + before the `post_list_feature_groups_with_metadata` interceptor. """ return response + async def post_list_feature_groups_with_metadata( + self, + response: feature_registry_service.ListFeatureGroupsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureGroupsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_groups + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_groups_with_metadata` + interceptor in new development instead of the `post_list_feature_groups` interceptor. + When both interceptors are used, this `post_list_feature_groups_with_metadata` interceptor runs after the + `post_list_feature_groups` interceptor. The (possibly modified) response returned by + `post_list_feature_groups` will be passed to + `post_list_feature_groups_with_metadata`. + """ + return response, metadata + async def pre_list_feature_monitor_jobs( self, request: feature_registry_service.ListFeatureMonitorJobsRequest, @@ -618,12 +925,38 @@ async def post_list_feature_monitor_jobs( ) -> feature_registry_service.ListFeatureMonitorJobsResponse: """Post-rpc interceptor for list_feature_monitor_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_monitor_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_monitor_jobs` interceptor runs + before the `post_list_feature_monitor_jobs_with_metadata` interceptor. """ return response + async def post_list_feature_monitor_jobs_with_metadata( + self, + response: feature_registry_service.ListFeatureMonitorJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureMonitorJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_monitor_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_monitor_jobs_with_metadata` + interceptor in new development instead of the `post_list_feature_monitor_jobs` interceptor. + When both interceptors are used, this `post_list_feature_monitor_jobs_with_metadata` interceptor runs after the + `post_list_feature_monitor_jobs` interceptor. The (possibly modified) response returned by + `post_list_feature_monitor_jobs` will be passed to + `post_list_feature_monitor_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_feature_monitors( self, request: feature_registry_service.ListFeatureMonitorsRequest, @@ -644,12 +977,38 @@ async def post_list_feature_monitors( ) -> feature_registry_service.ListFeatureMonitorsResponse: """Post-rpc interceptor for list_feature_monitors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_feature_monitors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_feature_monitors` interceptor runs + before the `post_list_feature_monitors_with_metadata` interceptor. """ return response + async def post_list_feature_monitors_with_metadata( + self, + response: feature_registry_service.ListFeatureMonitorsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + feature_registry_service.ListFeatureMonitorsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_feature_monitors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_feature_monitors_with_metadata` + interceptor in new development instead of the `post_list_feature_monitors` interceptor. + When both interceptors are used, this `post_list_feature_monitors_with_metadata` interceptor runs after the + `post_list_feature_monitors` interceptor. The (possibly modified) response returned by + `post_list_feature_monitors` will be passed to + `post_list_feature_monitors_with_metadata`. + """ + return response, metadata + async def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -670,12 +1029,38 @@ async def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + async def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + async def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -696,12 +1081,35 @@ async def post_update_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + async def post_update_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + async def pre_update_feature_group( self, request: feature_registry_service.UpdateFeatureGroupRequest, @@ -722,12 +1130,35 @@ async def post_update_feature_group( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_group - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_group_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_group` interceptor runs + before the `post_update_feature_group_with_metadata` interceptor. """ return response + async def post_update_feature_group_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_group + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_group_with_metadata` + interceptor in new development instead of the `post_update_feature_group` interceptor. + When both interceptors are used, this `post_update_feature_group_with_metadata` interceptor runs after the + `post_update_feature_group` interceptor. The (possibly modified) response returned by + `post_update_feature_group` will be passed to + `post_update_feature_group_with_metadata`. + """ + return response, metadata + async def pre_update_feature_monitor( self, request: feature_registry_service.UpdateFeatureMonitorRequest, @@ -748,12 +1179,35 @@ async def post_update_feature_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_feature_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeatureRegistryService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature_monitor` interceptor runs + before the `post_update_feature_monitor_with_metadata` interceptor. """ return response + async def post_update_feature_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeatureRegistryService server but before it is returned to user code. + + We recommend only using this `post_update_feature_monitor_with_metadata` + interceptor in new development instead of the `post_update_feature_monitor` interceptor. + When both interceptors are used, this `post_update_feature_monitor_with_metadata` interceptor runs after the + `post_update_feature_monitor` interceptor. The (possibly modified) response returned by + `post_update_feature_monitor` will be passed to + `post_update_feature_monitor_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1362,6 +1816,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1524,6 +1982,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1684,6 +2146,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1844,6 +2310,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2004,6 +2474,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature_monitor_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_feature_monitor_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2162,6 +2639,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2316,6 +2797,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2470,6 +2955,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2626,6 +3115,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2777,6 +3270,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2928,6 +3425,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3079,6 +3580,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature_monitor_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_feature_monitor_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3234,6 +3742,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_groups(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_feature_groups_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3393,6 +3905,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_monitor_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_feature_monitor_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3550,6 +4069,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_feature_monitors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_feature_monitors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3711,6 +4234,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3875,6 +4402,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4035,6 +4566,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_group(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_group_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4195,6 +4730,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/client.py b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/client.py index c822680129..eb3a0f5949 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -516,6 +518,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1192,16 +1221,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1247,16 +1280,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1419,16 +1456,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1540,16 +1581,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1662,16 +1707,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1722,16 +1771,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1777,16 +1830,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1832,16 +1889,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest.py index 39b9dd690f..14b8ccf3cf 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest.py @@ -128,12 +128,38 @@ def post_read_feature_values( ) -> featurestore_online_service.ReadFeatureValuesResponse: """Post-rpc interceptor for read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_read_feature_values` interceptor runs + before the `post_read_feature_values_with_metadata` interceptor. """ return response + def post_read_feature_values_with_metadata( + self, + response: featurestore_online_service.ReadFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.ReadFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_read_feature_values_with_metadata` + interceptor in new development instead of the `post_read_feature_values` interceptor. + When both interceptors are used, this `post_read_feature_values_with_metadata` interceptor runs after the + `post_read_feature_values` interceptor. The (possibly modified) response returned by + `post_read_feature_values` will be passed to + `post_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_streaming_read_feature_values( self, request: featurestore_online_service.StreamingReadFeatureValuesRequest, @@ -154,12 +180,37 @@ def post_streaming_read_feature_values( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for streaming_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_streaming_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_streaming_read_feature_values` interceptor runs + before the `post_streaming_read_feature_values_with_metadata` interceptor. """ return response + def post_streaming_read_feature_values_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for streaming_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_streaming_read_feature_values_with_metadata` + interceptor in new development instead of the `post_streaming_read_feature_values` interceptor. + When both interceptors are used, this `post_streaming_read_feature_values_with_metadata` interceptor runs after the + `post_streaming_read_feature_values` interceptor. The (possibly modified) response returned by + `post_streaming_read_feature_values` will be passed to + `post_streaming_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_write_feature_values( self, request: featurestore_online_service.WriteFeatureValuesRequest, @@ -180,12 +231,38 @@ def post_write_feature_values( ) -> featurestore_online_service.WriteFeatureValuesResponse: """Post-rpc interceptor for write_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_write_feature_values` interceptor runs + before the `post_write_feature_values_with_metadata` interceptor. """ return response + def post_write_feature_values_with_metadata( + self, + response: featurestore_online_service.WriteFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.WriteFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_write_feature_values_with_metadata` + interceptor in new development instead of the `post_write_feature_values` interceptor. + When both interceptors are used, this `post_write_feature_values_with_metadata` interceptor runs after the + `post_write_feature_values` interceptor. The (possibly modified) response returned by + `post_write_feature_values` will be passed to + `post_write_feature_values_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -652,6 +729,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -812,6 +893,13 @@ def __call__( ) resp = self._interceptor.post_streaming_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_streaming_read_feature_values_with_metadata( + resp, response_metadata + ) return resp class _WriteFeatureValues( @@ -946,6 +1034,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_write_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest_asyncio.py index 8ca29054f7..eb2ee5a527 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_online_serving_service/transports/rest_asyncio.py @@ -145,12 +145,38 @@ async def post_read_feature_values( ) -> featurestore_online_service.ReadFeatureValuesResponse: """Post-rpc interceptor for read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_read_feature_values` interceptor runs + before the `post_read_feature_values_with_metadata` interceptor. """ return response + async def post_read_feature_values_with_metadata( + self, + response: featurestore_online_service.ReadFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.ReadFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_read_feature_values_with_metadata` + interceptor in new development instead of the `post_read_feature_values` interceptor. + When both interceptors are used, this `post_read_feature_values_with_metadata` interceptor runs after the + `post_read_feature_values` interceptor. The (possibly modified) response returned by + `post_read_feature_values` will be passed to + `post_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_streaming_read_feature_values( self, request: featurestore_online_service.StreamingReadFeatureValuesRequest, @@ -171,12 +197,38 @@ async def post_streaming_read_feature_values( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for streaming_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_streaming_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_streaming_read_feature_values` interceptor runs + before the `post_streaming_read_feature_values_with_metadata` interceptor. """ return response + async def post_streaming_read_feature_values_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for streaming_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_streaming_read_feature_values_with_metadata` + interceptor in new development instead of the `post_streaming_read_feature_values` interceptor. + When both interceptors are used, this `post_streaming_read_feature_values_with_metadata` interceptor runs after the + `post_streaming_read_feature_values` interceptor. The (possibly modified) response returned by + `post_streaming_read_feature_values` will be passed to + `post_streaming_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_write_feature_values( self, request: featurestore_online_service.WriteFeatureValuesRequest, @@ -197,12 +249,38 @@ async def post_write_feature_values( ) -> featurestore_online_service.WriteFeatureValuesResponse: """Post-rpc interceptor for write_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreOnlineServingService server but before - it is returned to user code. + it is returned to user code. This `post_write_feature_values` interceptor runs + before the `post_write_feature_values_with_metadata` interceptor. """ return response + async def post_write_feature_values_with_metadata( + self, + response: featurestore_online_service.WriteFeatureValuesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_online_service.WriteFeatureValuesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreOnlineServingService server but before it is returned to user code. + + We recommend only using this `post_write_feature_values_with_metadata` + interceptor in new development instead of the `post_write_feature_values` interceptor. + When both interceptors are used, this `post_write_feature_values_with_metadata` interceptor runs after the + `post_write_feature_values` interceptor. The (possibly modified) response returned by + `post_write_feature_values` will be passed to + `post_write_feature_values_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -730,6 +808,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -898,6 +980,13 @@ async def __call__( response, featurestore_online_service.ReadFeatureValuesResponse ) resp = await self._interceptor.post_streaming_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_streaming_read_feature_values_with_metadata( + resp, response_metadata + ) return resp class _WriteFeatureValues( @@ -1037,6 +1126,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_write_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_service/client.py b/google/cloud/aiplatform_v1beta1/services/featurestore_service/client.py index eb21ddee64..7e084785f4 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -569,6 +571,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3790,16 +3819,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3845,16 +3878,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -4017,16 +4054,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -4138,16 +4179,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -4260,16 +4305,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -4320,16 +4369,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -4375,16 +4428,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -4430,16 +4487,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest.py index aeab4dab5b..017480b816 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest.py @@ -278,12 +278,35 @@ def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + def pre_batch_read_feature_values( self, request: featurestore_service.BatchReadFeatureValuesRequest, @@ -304,12 +327,35 @@ def post_batch_read_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_feature_values` interceptor runs + before the `post_batch_read_feature_values_with_metadata` interceptor. """ return response + def post_batch_read_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_read_feature_values_with_metadata` + interceptor in new development instead of the `post_batch_read_feature_values` interceptor. + When both interceptors are used, this `post_batch_read_feature_values_with_metadata` interceptor runs after the + `post_batch_read_feature_values` interceptor. The (possibly modified) response returned by + `post_batch_read_feature_values` will be passed to + `post_batch_read_feature_values_with_metadata`. + """ + return response, metadata + def pre_create_entity_type( self, request: featurestore_service.CreateEntityTypeRequest, @@ -330,12 +376,35 @@ def post_create_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_entity_type` interceptor runs + before the `post_create_entity_type_with_metadata` interceptor. """ return response + def post_create_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_entity_type_with_metadata` + interceptor in new development instead of the `post_create_entity_type` interceptor. + When both interceptors are used, this `post_create_entity_type_with_metadata` interceptor runs after the + `post_create_entity_type` interceptor. The (possibly modified) response returned by + `post_create_entity_type` will be passed to + `post_create_entity_type_with_metadata`. + """ + return response, metadata + def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -356,12 +425,35 @@ def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + def pre_create_featurestore( self, request: featurestore_service.CreateFeaturestoreRequest, @@ -382,12 +474,35 @@ def post_create_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_featurestore` interceptor runs + before the `post_create_featurestore_with_metadata` interceptor. """ return response + def post_create_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_featurestore_with_metadata` + interceptor in new development instead of the `post_create_featurestore` interceptor. + When both interceptors are used, this `post_create_featurestore_with_metadata` interceptor runs after the + `post_create_featurestore` interceptor. The (possibly modified) response returned by + `post_create_featurestore` will be passed to + `post_create_featurestore_with_metadata`. + """ + return response, metadata + def pre_delete_entity_type( self, request: featurestore_service.DeleteEntityTypeRequest, @@ -408,12 +523,35 @@ def post_delete_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_entity_type` interceptor runs + before the `post_delete_entity_type_with_metadata` interceptor. """ return response + def post_delete_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_entity_type_with_metadata` + interceptor in new development instead of the `post_delete_entity_type` interceptor. + When both interceptors are used, this `post_delete_entity_type_with_metadata` interceptor runs after the + `post_delete_entity_type` interceptor. The (possibly modified) response returned by + `post_delete_entity_type` will be passed to + `post_delete_entity_type_with_metadata`. + """ + return response, metadata + def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -434,12 +572,35 @@ def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + def pre_delete_featurestore( self, request: featurestore_service.DeleteFeaturestoreRequest, @@ -460,12 +621,35 @@ def post_delete_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_featurestore` interceptor runs + before the `post_delete_featurestore_with_metadata` interceptor. """ return response + def post_delete_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_featurestore_with_metadata` + interceptor in new development instead of the `post_delete_featurestore` interceptor. + When both interceptors are used, this `post_delete_featurestore_with_metadata` interceptor runs after the + `post_delete_featurestore` interceptor. The (possibly modified) response returned by + `post_delete_featurestore` will be passed to + `post_delete_featurestore_with_metadata`. + """ + return response, metadata + def pre_delete_feature_values( self, request: featurestore_service.DeleteFeatureValuesRequest, @@ -486,12 +670,35 @@ def post_delete_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_values` interceptor runs + before the `post_delete_feature_values_with_metadata` interceptor. """ return response + def post_delete_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_values_with_metadata` + interceptor in new development instead of the `post_delete_feature_values` interceptor. + When both interceptors are used, this `post_delete_feature_values_with_metadata` interceptor runs after the + `post_delete_feature_values` interceptor. The (possibly modified) response returned by + `post_delete_feature_values` will be passed to + `post_delete_feature_values_with_metadata`. + """ + return response, metadata + def pre_export_feature_values( self, request: featurestore_service.ExportFeatureValuesRequest, @@ -512,12 +719,35 @@ def post_export_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_export_feature_values` interceptor runs + before the `post_export_feature_values_with_metadata` interceptor. """ return response + def post_export_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_export_feature_values_with_metadata` + interceptor in new development instead of the `post_export_feature_values` interceptor. + When both interceptors are used, this `post_export_feature_values_with_metadata` interceptor runs after the + `post_export_feature_values` interceptor. The (possibly modified) response returned by + `post_export_feature_values` will be passed to + `post_export_feature_values_with_metadata`. + """ + return response, metadata + def pre_get_entity_type( self, request: featurestore_service.GetEntityTypeRequest, @@ -538,12 +768,35 @@ def post_get_entity_type( ) -> entity_type.EntityType: """Post-rpc interceptor for get_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_entity_type` interceptor runs + before the `post_get_entity_type_with_metadata` interceptor. """ return response + def post_get_entity_type_with_metadata( + self, + response: entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_entity_type_with_metadata` + interceptor in new development instead of the `post_get_entity_type` interceptor. + When both interceptors are used, this `post_get_entity_type_with_metadata` interceptor runs after the + `post_get_entity_type` interceptor. The (possibly modified) response returned by + `post_get_entity_type` will be passed to + `post_get_entity_type_with_metadata`. + """ + return response, metadata + def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -561,12 +814,35 @@ def pre_get_feature( def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + def pre_get_featurestore( self, request: featurestore_service.GetFeaturestoreRequest, @@ -587,12 +863,35 @@ def post_get_featurestore( ) -> featurestore.Featurestore: """Post-rpc interceptor for get_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_featurestore` interceptor runs + before the `post_get_featurestore_with_metadata` interceptor. """ return response + def post_get_featurestore_with_metadata( + self, + response: featurestore.Featurestore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[featurestore.Featurestore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_featurestore_with_metadata` + interceptor in new development instead of the `post_get_featurestore` interceptor. + When both interceptors are used, this `post_get_featurestore_with_metadata` interceptor runs after the + `post_get_featurestore` interceptor. The (possibly modified) response returned by + `post_get_featurestore` will be passed to + `post_get_featurestore_with_metadata`. + """ + return response, metadata + def pre_import_feature_values( self, request: featurestore_service.ImportFeatureValuesRequest, @@ -613,12 +912,35 @@ def post_import_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_import_feature_values` interceptor runs + before the `post_import_feature_values_with_metadata` interceptor. """ return response + def post_import_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_import_feature_values_with_metadata` + interceptor in new development instead of the `post_import_feature_values` interceptor. + When both interceptors are used, this `post_import_feature_values_with_metadata` interceptor runs after the + `post_import_feature_values` interceptor. The (possibly modified) response returned by + `post_import_feature_values` will be passed to + `post_import_feature_values_with_metadata`. + """ + return response, metadata + def pre_list_entity_types( self, request: featurestore_service.ListEntityTypesRequest, @@ -639,12 +961,38 @@ def post_list_entity_types( ) -> featurestore_service.ListEntityTypesResponse: """Post-rpc interceptor for list_entity_types - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_entity_types_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_entity_types` interceptor runs + before the `post_list_entity_types_with_metadata` interceptor. """ return response + def post_list_entity_types_with_metadata( + self, + response: featurestore_service.ListEntityTypesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListEntityTypesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_entity_types + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_entity_types_with_metadata` + interceptor in new development instead of the `post_list_entity_types` interceptor. + When both interceptors are used, this `post_list_entity_types_with_metadata` interceptor runs after the + `post_list_entity_types` interceptor. The (possibly modified) response returned by + `post_list_entity_types` will be passed to + `post_list_entity_types_with_metadata`. + """ + return response, metadata + def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -665,12 +1013,38 @@ def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + def pre_list_featurestores( self, request: featurestore_service.ListFeaturestoresRequest, @@ -691,12 +1065,38 @@ def post_list_featurestores( ) -> featurestore_service.ListFeaturestoresResponse: """Post-rpc interceptor for list_featurestores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_featurestores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_featurestores` interceptor runs + before the `post_list_featurestores_with_metadata` interceptor. """ return response + def post_list_featurestores_with_metadata( + self, + response: featurestore_service.ListFeaturestoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturestoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_featurestores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_featurestores_with_metadata` + interceptor in new development instead of the `post_list_featurestores` interceptor. + When both interceptors are used, this `post_list_featurestores_with_metadata` interceptor runs after the + `post_list_featurestores` interceptor. The (possibly modified) response returned by + `post_list_featurestores` will be passed to + `post_list_featurestores_with_metadata`. + """ + return response, metadata + def pre_search_features( self, request: featurestore_service.SearchFeaturesRequest, @@ -717,12 +1117,38 @@ def post_search_features( ) -> featurestore_service.SearchFeaturesResponse: """Post-rpc interceptor for search_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_features` interceptor runs + before the `post_search_features_with_metadata` interceptor. """ return response + def post_search_features_with_metadata( + self, + response: featurestore_service.SearchFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.SearchFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_search_features_with_metadata` + interceptor in new development instead of the `post_search_features` interceptor. + When both interceptors are used, this `post_search_features_with_metadata` interceptor runs after the + `post_search_features` interceptor. The (possibly modified) response returned by + `post_search_features` will be passed to + `post_search_features_with_metadata`. + """ + return response, metadata + def pre_update_entity_type( self, request: featurestore_service.UpdateEntityTypeRequest, @@ -743,12 +1169,35 @@ def post_update_entity_type( ) -> gca_entity_type.EntityType: """Post-rpc interceptor for update_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_entity_type` interceptor runs + before the `post_update_entity_type_with_metadata` interceptor. """ return response + def post_update_entity_type_with_metadata( + self, + response: gca_entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_entity_type_with_metadata` + interceptor in new development instead of the `post_update_entity_type` interceptor. + When both interceptors are used, this `post_update_entity_type_with_metadata` interceptor runs after the + `post_update_entity_type` interceptor. The (possibly modified) response returned by + `post_update_entity_type` will be passed to + `post_update_entity_type_with_metadata`. + """ + return response, metadata + def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -767,12 +1216,35 @@ def pre_update_feature( def post_update_feature(self, response: gca_feature.Feature) -> gca_feature.Feature: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + def post_update_feature_with_metadata( + self, + response: gca_feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + def pre_update_featurestore( self, request: featurestore_service.UpdateFeaturestoreRequest, @@ -793,12 +1265,35 @@ def post_update_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_featurestore` interceptor runs + before the `post_update_featurestore_with_metadata` interceptor. """ return response + def post_update_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_featurestore_with_metadata` + interceptor in new development instead of the `post_update_featurestore` interceptor. + When both interceptors are used, this `post_update_featurestore_with_metadata` interceptor runs after the + `post_update_featurestore` interceptor. The (possibly modified) response returned by + `post_update_featurestore` will be passed to + `post_update_featurestore_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3350,6 +3845,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3504,6 +4003,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3656,6 +4159,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3808,6 +4315,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3962,6 +4473,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4108,6 +4623,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4254,6 +4773,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4402,6 +4925,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4556,6 +5083,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4710,6 +5241,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4859,6 +5394,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5007,6 +5546,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5157,6 +5700,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5311,6 +5858,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5458,6 +6009,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_entity_types(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_entity_types_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5609,6 +6164,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5760,6 +6319,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_featurestores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_featurestores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5907,6 +6470,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6066,6 +6633,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6220,6 +6791,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6374,6 +6949,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest_asyncio.py index cca10648c7..685ca09ff6 100644 --- a/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/featurestore_service/transports/rest_asyncio.py @@ -296,12 +296,35 @@ async def post_batch_create_features( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_create_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_features` interceptor runs + before the `post_batch_create_features_with_metadata` interceptor. """ return response + async def post_batch_create_features_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_create_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_create_features_with_metadata` + interceptor in new development instead of the `post_batch_create_features` interceptor. + When both interceptors are used, this `post_batch_create_features_with_metadata` interceptor runs after the + `post_batch_create_features` interceptor. The (possibly modified) response returned by + `post_batch_create_features` will be passed to + `post_batch_create_features_with_metadata`. + """ + return response, metadata + async def pre_batch_read_feature_values( self, request: featurestore_service.BatchReadFeatureValuesRequest, @@ -322,12 +345,35 @@ async def post_batch_read_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_read_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_feature_values` interceptor runs + before the `post_batch_read_feature_values_with_metadata` interceptor. """ return response + async def post_batch_read_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_read_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_batch_read_feature_values_with_metadata` + interceptor in new development instead of the `post_batch_read_feature_values` interceptor. + When both interceptors are used, this `post_batch_read_feature_values_with_metadata` interceptor runs after the + `post_batch_read_feature_values` interceptor. The (possibly modified) response returned by + `post_batch_read_feature_values` will be passed to + `post_batch_read_feature_values_with_metadata`. + """ + return response, metadata + async def pre_create_entity_type( self, request: featurestore_service.CreateEntityTypeRequest, @@ -348,12 +394,35 @@ async def post_create_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_entity_type` interceptor runs + before the `post_create_entity_type_with_metadata` interceptor. """ return response + async def post_create_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_entity_type_with_metadata` + interceptor in new development instead of the `post_create_entity_type` interceptor. + When both interceptors are used, this `post_create_entity_type_with_metadata` interceptor runs after the + `post_create_entity_type` interceptor. The (possibly modified) response returned by + `post_create_entity_type` will be passed to + `post_create_entity_type_with_metadata`. + """ + return response, metadata + async def pre_create_feature( self, request: featurestore_service.CreateFeatureRequest, @@ -374,12 +443,35 @@ async def post_create_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_feature` interceptor runs + before the `post_create_feature_with_metadata` interceptor. """ return response + async def post_create_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_feature_with_metadata` + interceptor in new development instead of the `post_create_feature` interceptor. + When both interceptors are used, this `post_create_feature_with_metadata` interceptor runs after the + `post_create_feature` interceptor. The (possibly modified) response returned by + `post_create_feature` will be passed to + `post_create_feature_with_metadata`. + """ + return response, metadata + async def pre_create_featurestore( self, request: featurestore_service.CreateFeaturestoreRequest, @@ -400,12 +492,35 @@ async def post_create_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_create_featurestore` interceptor runs + before the `post_create_featurestore_with_metadata` interceptor. """ return response + async def post_create_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_create_featurestore_with_metadata` + interceptor in new development instead of the `post_create_featurestore` interceptor. + When both interceptors are used, this `post_create_featurestore_with_metadata` interceptor runs after the + `post_create_featurestore` interceptor. The (possibly modified) response returned by + `post_create_featurestore` will be passed to + `post_create_featurestore_with_metadata`. + """ + return response, metadata + async def pre_delete_entity_type( self, request: featurestore_service.DeleteEntityTypeRequest, @@ -426,12 +541,35 @@ async def post_delete_entity_type( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_entity_type` interceptor runs + before the `post_delete_entity_type_with_metadata` interceptor. """ return response + async def post_delete_entity_type_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_entity_type_with_metadata` + interceptor in new development instead of the `post_delete_entity_type` interceptor. + When both interceptors are used, this `post_delete_entity_type_with_metadata` interceptor runs after the + `post_delete_entity_type` interceptor. The (possibly modified) response returned by + `post_delete_entity_type` will be passed to + `post_delete_entity_type_with_metadata`. + """ + return response, metadata + async def pre_delete_feature( self, request: featurestore_service.DeleteFeatureRequest, @@ -452,12 +590,35 @@ async def post_delete_feature( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature` interceptor runs + before the `post_delete_feature_with_metadata` interceptor. """ return response + async def post_delete_feature_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_with_metadata` + interceptor in new development instead of the `post_delete_feature` interceptor. + When both interceptors are used, this `post_delete_feature_with_metadata` interceptor runs after the + `post_delete_feature` interceptor. The (possibly modified) response returned by + `post_delete_feature` will be passed to + `post_delete_feature_with_metadata`. + """ + return response, metadata + async def pre_delete_featurestore( self, request: featurestore_service.DeleteFeaturestoreRequest, @@ -478,12 +639,35 @@ async def post_delete_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_featurestore` interceptor runs + before the `post_delete_featurestore_with_metadata` interceptor. """ return response + async def post_delete_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_featurestore_with_metadata` + interceptor in new development instead of the `post_delete_featurestore` interceptor. + When both interceptors are used, this `post_delete_featurestore_with_metadata` interceptor runs after the + `post_delete_featurestore` interceptor. The (possibly modified) response returned by + `post_delete_featurestore` will be passed to + `post_delete_featurestore_with_metadata`. + """ + return response, metadata + async def pre_delete_feature_values( self, request: featurestore_service.DeleteFeatureValuesRequest, @@ -504,12 +688,35 @@ async def post_delete_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_delete_feature_values` interceptor runs + before the `post_delete_feature_values_with_metadata` interceptor. """ return response + async def post_delete_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_delete_feature_values_with_metadata` + interceptor in new development instead of the `post_delete_feature_values` interceptor. + When both interceptors are used, this `post_delete_feature_values_with_metadata` interceptor runs after the + `post_delete_feature_values` interceptor. The (possibly modified) response returned by + `post_delete_feature_values` will be passed to + `post_delete_feature_values_with_metadata`. + """ + return response, metadata + async def pre_export_feature_values( self, request: featurestore_service.ExportFeatureValuesRequest, @@ -530,12 +737,35 @@ async def post_export_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_export_feature_values` interceptor runs + before the `post_export_feature_values_with_metadata` interceptor. """ return response + async def post_export_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_export_feature_values_with_metadata` + interceptor in new development instead of the `post_export_feature_values` interceptor. + When both interceptors are used, this `post_export_feature_values_with_metadata` interceptor runs after the + `post_export_feature_values` interceptor. The (possibly modified) response returned by + `post_export_feature_values` will be passed to + `post_export_feature_values_with_metadata`. + """ + return response, metadata + async def pre_get_entity_type( self, request: featurestore_service.GetEntityTypeRequest, @@ -556,12 +786,35 @@ async def post_get_entity_type( ) -> entity_type.EntityType: """Post-rpc interceptor for get_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_entity_type` interceptor runs + before the `post_get_entity_type_with_metadata` interceptor. """ return response + async def post_get_entity_type_with_metadata( + self, + response: entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_entity_type_with_metadata` + interceptor in new development instead of the `post_get_entity_type` interceptor. + When both interceptors are used, this `post_get_entity_type_with_metadata` interceptor runs after the + `post_get_entity_type` interceptor. The (possibly modified) response returned by + `post_get_entity_type` will be passed to + `post_get_entity_type_with_metadata`. + """ + return response, metadata + async def pre_get_feature( self, request: featurestore_service.GetFeatureRequest, @@ -579,12 +832,35 @@ async def pre_get_feature( async def post_get_feature(self, response: feature.Feature) -> feature.Feature: """Post-rpc interceptor for get_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_feature` interceptor runs + before the `post_get_feature_with_metadata` interceptor. """ return response + async def post_get_feature_with_metadata( + self, + response: feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_feature_with_metadata` + interceptor in new development instead of the `post_get_feature` interceptor. + When both interceptors are used, this `post_get_feature_with_metadata` interceptor runs after the + `post_get_feature` interceptor. The (possibly modified) response returned by + `post_get_feature` will be passed to + `post_get_feature_with_metadata`. + """ + return response, metadata + async def pre_get_featurestore( self, request: featurestore_service.GetFeaturestoreRequest, @@ -605,12 +881,35 @@ async def post_get_featurestore( ) -> featurestore.Featurestore: """Post-rpc interceptor for get_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_get_featurestore` interceptor runs + before the `post_get_featurestore_with_metadata` interceptor. """ return response + async def post_get_featurestore_with_metadata( + self, + response: featurestore.Featurestore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[featurestore.Featurestore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_get_featurestore_with_metadata` + interceptor in new development instead of the `post_get_featurestore` interceptor. + When both interceptors are used, this `post_get_featurestore_with_metadata` interceptor runs after the + `post_get_featurestore` interceptor. The (possibly modified) response returned by + `post_get_featurestore` will be passed to + `post_get_featurestore_with_metadata`. + """ + return response, metadata + async def pre_import_feature_values( self, request: featurestore_service.ImportFeatureValuesRequest, @@ -631,12 +930,35 @@ async def post_import_feature_values( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_feature_values - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_feature_values_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_import_feature_values` interceptor runs + before the `post_import_feature_values_with_metadata` interceptor. """ return response + async def post_import_feature_values_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_feature_values + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_import_feature_values_with_metadata` + interceptor in new development instead of the `post_import_feature_values` interceptor. + When both interceptors are used, this `post_import_feature_values_with_metadata` interceptor runs after the + `post_import_feature_values` interceptor. The (possibly modified) response returned by + `post_import_feature_values` will be passed to + `post_import_feature_values_with_metadata`. + """ + return response, metadata + async def pre_list_entity_types( self, request: featurestore_service.ListEntityTypesRequest, @@ -657,12 +979,38 @@ async def post_list_entity_types( ) -> featurestore_service.ListEntityTypesResponse: """Post-rpc interceptor for list_entity_types - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_entity_types_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_entity_types` interceptor runs + before the `post_list_entity_types_with_metadata` interceptor. """ return response + async def post_list_entity_types_with_metadata( + self, + response: featurestore_service.ListEntityTypesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListEntityTypesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_entity_types + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_entity_types_with_metadata` + interceptor in new development instead of the `post_list_entity_types` interceptor. + When both interceptors are used, this `post_list_entity_types_with_metadata` interceptor runs after the + `post_list_entity_types` interceptor. The (possibly modified) response returned by + `post_list_entity_types` will be passed to + `post_list_entity_types_with_metadata`. + """ + return response, metadata + async def pre_list_features( self, request: featurestore_service.ListFeaturesRequest, @@ -683,12 +1031,38 @@ async def post_list_features( ) -> featurestore_service.ListFeaturesResponse: """Post-rpc interceptor for list_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_features` interceptor runs + before the `post_list_features_with_metadata` interceptor. """ return response + async def post_list_features_with_metadata( + self, + response: featurestore_service.ListFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_features_with_metadata` + interceptor in new development instead of the `post_list_features` interceptor. + When both interceptors are used, this `post_list_features_with_metadata` interceptor runs after the + `post_list_features` interceptor. The (possibly modified) response returned by + `post_list_features` will be passed to + `post_list_features_with_metadata`. + """ + return response, metadata + async def pre_list_featurestores( self, request: featurestore_service.ListFeaturestoresRequest, @@ -709,12 +1083,38 @@ async def post_list_featurestores( ) -> featurestore_service.ListFeaturestoresResponse: """Post-rpc interceptor for list_featurestores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_featurestores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_list_featurestores` interceptor runs + before the `post_list_featurestores_with_metadata` interceptor. """ return response + async def post_list_featurestores_with_metadata( + self, + response: featurestore_service.ListFeaturestoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.ListFeaturestoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_featurestores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_list_featurestores_with_metadata` + interceptor in new development instead of the `post_list_featurestores` interceptor. + When both interceptors are used, this `post_list_featurestores_with_metadata` interceptor runs after the + `post_list_featurestores` interceptor. The (possibly modified) response returned by + `post_list_featurestores` will be passed to + `post_list_featurestores_with_metadata`. + """ + return response, metadata + async def pre_search_features( self, request: featurestore_service.SearchFeaturesRequest, @@ -735,12 +1135,38 @@ async def post_search_features( ) -> featurestore_service.SearchFeaturesResponse: """Post-rpc interceptor for search_features - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_features_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_search_features` interceptor runs + before the `post_search_features_with_metadata` interceptor. """ return response + async def post_search_features_with_metadata( + self, + response: featurestore_service.SearchFeaturesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + featurestore_service.SearchFeaturesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_features + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_search_features_with_metadata` + interceptor in new development instead of the `post_search_features` interceptor. + When both interceptors are used, this `post_search_features_with_metadata` interceptor runs after the + `post_search_features` interceptor. The (possibly modified) response returned by + `post_search_features` will be passed to + `post_search_features_with_metadata`. + """ + return response, metadata + async def pre_update_entity_type( self, request: featurestore_service.UpdateEntityTypeRequest, @@ -761,12 +1187,35 @@ async def post_update_entity_type( ) -> gca_entity_type.EntityType: """Post-rpc interceptor for update_entity_type - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_entity_type_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_entity_type` interceptor runs + before the `post_update_entity_type_with_metadata` interceptor. """ return response + async def post_update_entity_type_with_metadata( + self, + response: gca_entity_type.EntityType, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_entity_type.EntityType, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_entity_type + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_entity_type_with_metadata` + interceptor in new development instead of the `post_update_entity_type` interceptor. + When both interceptors are used, this `post_update_entity_type_with_metadata` interceptor runs after the + `post_update_entity_type` interceptor. The (possibly modified) response returned by + `post_update_entity_type` will be passed to + `post_update_entity_type_with_metadata`. + """ + return response, metadata + async def pre_update_feature( self, request: featurestore_service.UpdateFeatureRequest, @@ -787,12 +1236,35 @@ async def post_update_feature( ) -> gca_feature.Feature: """Post-rpc interceptor for update_feature - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_feature_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_feature` interceptor runs + before the `post_update_feature_with_metadata` interceptor. """ return response + async def post_update_feature_with_metadata( + self, + response: gca_feature.Feature, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_feature.Feature, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_feature + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_feature_with_metadata` + interceptor in new development instead of the `post_update_feature` interceptor. + When both interceptors are used, this `post_update_feature_with_metadata` interceptor runs after the + `post_update_feature` interceptor. The (possibly modified) response returned by + `post_update_feature` will be passed to + `post_update_feature_with_metadata`. + """ + return response, metadata + async def pre_update_featurestore( self, request: featurestore_service.UpdateFeaturestoreRequest, @@ -813,12 +1285,35 @@ async def post_update_featurestore( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_featurestore - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_featurestore_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the FeaturestoreService server but before - it is returned to user code. + it is returned to user code. This `post_update_featurestore` interceptor runs + before the `post_update_featurestore_with_metadata` interceptor. """ return response + async def post_update_featurestore_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_featurestore + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the FeaturestoreService server but before it is returned to user code. + + We recommend only using this `post_update_featurestore_with_metadata` + interceptor in new development instead of the `post_update_featurestore` interceptor. + When both interceptors are used, this `post_update_featurestore_with_metadata` interceptor runs after the + `post_update_featurestore` interceptor. The (possibly modified) response returned by + `post_update_featurestore` will be passed to + `post_update_featurestore_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1435,6 +1930,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_batch_create_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1595,6 +2094,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_read_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_read_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1755,6 +2261,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1917,6 +2427,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2077,6 +2591,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2231,6 +2749,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2387,6 +2909,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2541,6 +3067,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2701,6 +3231,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2861,6 +3395,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3018,6 +3556,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3176,6 +3718,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3332,6 +3878,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3492,6 +4042,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_feature_values(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_feature_values_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3645,6 +4199,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_entity_types(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_entity_types_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3806,6 +4364,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3961,6 +4523,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_featurestores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_featurestores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4116,6 +4682,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_features(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_search_features_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4281,6 +4851,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_entity_type(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_entity_type_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4443,6 +5017,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_feature(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_feature_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4603,6 +5181,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_featurestore(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_featurestore_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/client.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/client.py index 689426bfd4..d60d95ca4f 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -530,6 +532,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1377,16 +1406,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1432,16 +1465,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1604,16 +1641,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1725,16 +1766,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1847,16 +1892,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1907,16 +1956,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1962,16 +2015,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2017,16 +2074,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest.py index 0c3a0f41f2..1a4ba65a6b 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest.py @@ -143,12 +143,37 @@ def post_create_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for create_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_create_cached_content` interceptor runs + before the `post_create_cached_content_with_metadata` interceptor. """ return response + def post_create_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_create_cached_content_with_metadata` + interceptor in new development instead of the `post_create_cached_content` interceptor. + When both interceptors are used, this `post_create_cached_content_with_metadata` interceptor runs after the + `post_create_cached_content` interceptor. The (possibly modified) response returned by + `post_create_cached_content` will be passed to + `post_create_cached_content_with_metadata`. + """ + return response, metadata + def pre_delete_cached_content( self, request: gen_ai_cache_service.DeleteCachedContentRequest, @@ -184,12 +209,35 @@ def post_get_cached_content( ) -> cached_content.CachedContent: """Post-rpc interceptor for get_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_get_cached_content` interceptor runs + before the `post_get_cached_content_with_metadata` interceptor. """ return response + def post_get_cached_content_with_metadata( + self, + response: cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_get_cached_content_with_metadata` + interceptor in new development instead of the `post_get_cached_content` interceptor. + When both interceptors are used, this `post_get_cached_content_with_metadata` interceptor runs after the + `post_get_cached_content` interceptor. The (possibly modified) response returned by + `post_get_cached_content` will be passed to + `post_get_cached_content_with_metadata`. + """ + return response, metadata + def pre_list_cached_contents( self, request: gen_ai_cache_service.ListCachedContentsRequest, @@ -210,12 +258,38 @@ def post_list_cached_contents( ) -> gen_ai_cache_service.ListCachedContentsResponse: """Post-rpc interceptor for list_cached_contents - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_cached_contents_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_list_cached_contents` interceptor runs + before the `post_list_cached_contents_with_metadata` interceptor. """ return response + def post_list_cached_contents_with_metadata( + self, + response: gen_ai_cache_service.ListCachedContentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gen_ai_cache_service.ListCachedContentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_cached_contents + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_list_cached_contents_with_metadata` + interceptor in new development instead of the `post_list_cached_contents` interceptor. + When both interceptors are used, this `post_list_cached_contents_with_metadata` interceptor runs after the + `post_list_cached_contents` interceptor. The (possibly modified) response returned by + `post_list_cached_contents` will be passed to + `post_list_cached_contents_with_metadata`. + """ + return response, metadata + def pre_update_cached_content( self, request: gen_ai_cache_service.UpdateCachedContentRequest, @@ -236,12 +310,37 @@ def post_update_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for update_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_update_cached_content` interceptor runs + before the `post_update_cached_content_with_metadata` interceptor. """ return response + def post_update_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_update_cached_content_with_metadata` + interceptor in new development instead of the `post_update_cached_content` interceptor. + When both interceptors are used, this `post_update_cached_content_with_metadata` interceptor runs after the + `post_update_cached_content` interceptor. The (possibly modified) response returned by + `post_update_cached_content` will be passed to + `post_update_cached_content_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -705,6 +804,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -968,6 +1071,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1114,6 +1221,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_cached_contents(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_cached_contents_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1275,6 +1386,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest_asyncio.py index dcb62371c5..6add2c0b17 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_cache_service/transports/rest_asyncio.py @@ -160,12 +160,37 @@ async def post_create_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for create_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_create_cached_content` interceptor runs + before the `post_create_cached_content_with_metadata` interceptor. """ return response + async def post_create_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_create_cached_content_with_metadata` + interceptor in new development instead of the `post_create_cached_content` interceptor. + When both interceptors are used, this `post_create_cached_content_with_metadata` interceptor runs after the + `post_create_cached_content` interceptor. The (possibly modified) response returned by + `post_create_cached_content` will be passed to + `post_create_cached_content_with_metadata`. + """ + return response, metadata + async def pre_delete_cached_content( self, request: gen_ai_cache_service.DeleteCachedContentRequest, @@ -201,12 +226,35 @@ async def post_get_cached_content( ) -> cached_content.CachedContent: """Post-rpc interceptor for get_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_get_cached_content` interceptor runs + before the `post_get_cached_content_with_metadata` interceptor. """ return response + async def post_get_cached_content_with_metadata( + self, + response: cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_get_cached_content_with_metadata` + interceptor in new development instead of the `post_get_cached_content` interceptor. + When both interceptors are used, this `post_get_cached_content_with_metadata` interceptor runs after the + `post_get_cached_content` interceptor. The (possibly modified) response returned by + `post_get_cached_content` will be passed to + `post_get_cached_content_with_metadata`. + """ + return response, metadata + async def pre_list_cached_contents( self, request: gen_ai_cache_service.ListCachedContentsRequest, @@ -227,12 +275,38 @@ async def post_list_cached_contents( ) -> gen_ai_cache_service.ListCachedContentsResponse: """Post-rpc interceptor for list_cached_contents - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_cached_contents_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_list_cached_contents` interceptor runs + before the `post_list_cached_contents_with_metadata` interceptor. """ return response + async def post_list_cached_contents_with_metadata( + self, + response: gen_ai_cache_service.ListCachedContentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gen_ai_cache_service.ListCachedContentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_cached_contents + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_list_cached_contents_with_metadata` + interceptor in new development instead of the `post_list_cached_contents` interceptor. + When both interceptors are used, this `post_list_cached_contents_with_metadata` interceptor runs after the + `post_list_cached_contents` interceptor. The (possibly modified) response returned by + `post_list_cached_contents` will be passed to + `post_list_cached_contents_with_metadata`. + """ + return response, metadata + async def pre_update_cached_content( self, request: gen_ai_cache_service.UpdateCachedContentRequest, @@ -253,12 +327,37 @@ async def post_update_cached_content( ) -> gca_cached_content.CachedContent: """Post-rpc interceptor for update_cached_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_cached_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiCacheService server but before - it is returned to user code. + it is returned to user code. This `post_update_cached_content` interceptor runs + before the `post_update_cached_content_with_metadata` interceptor. """ return response + async def post_update_cached_content_with_metadata( + self, + response: gca_cached_content.CachedContent, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_cached_content.CachedContent, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_cached_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiCacheService server but before it is returned to user code. + + We recommend only using this `post_update_cached_content_with_metadata` + interceptor in new development instead of the `post_update_cached_content` interceptor. + When both interceptors are used, this `post_update_cached_content_with_metadata` interceptor runs after the + `post_update_cached_content` interceptor. The (possibly modified) response returned by + `post_update_cached_content` will be passed to + `post_update_cached_content_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -789,6 +888,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1062,6 +1165,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1214,6 +1321,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_cached_contents(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_cached_contents_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1379,6 +1490,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_cached_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_cached_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/client.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/client.py index 4041061223..a31679b6d6 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -599,6 +601,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1480,16 +1509,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1535,16 +1568,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1707,16 +1744,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1828,16 +1869,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1950,16 +1995,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2010,16 +2059,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2065,16 +2118,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2120,16 +2177,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest.py index 625248bba9..c42df9d90a 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest.py @@ -159,12 +159,35 @@ def post_create_tuning_job( ) -> gca_tuning_job.TuningJob: """Post-rpc interceptor for create_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_create_tuning_job` interceptor runs + before the `post_create_tuning_job_with_metadata` interceptor. """ return response + def post_create_tuning_job_with_metadata( + self, + response: gca_tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_create_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_tuning_job` interceptor. + When both interceptors are used, this `post_create_tuning_job_with_metadata` interceptor runs after the + `post_create_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_tuning_job` will be passed to + `post_create_tuning_job_with_metadata`. + """ + return response, metadata + def pre_get_tuning_job( self, request: genai_tuning_service.GetTuningJobRequest, @@ -185,12 +208,35 @@ def post_get_tuning_job( ) -> tuning_job.TuningJob: """Post-rpc interceptor for get_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_get_tuning_job` interceptor runs + before the `post_get_tuning_job_with_metadata` interceptor. """ return response + def post_get_tuning_job_with_metadata( + self, + response: tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_get_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_tuning_job` interceptor. + When both interceptors are used, this `post_get_tuning_job_with_metadata` interceptor runs after the + `post_get_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_tuning_job` will be passed to + `post_get_tuning_job_with_metadata`. + """ + return response, metadata + def pre_list_tuning_jobs( self, request: genai_tuning_service.ListTuningJobsRequest, @@ -211,12 +257,38 @@ def post_list_tuning_jobs( ) -> genai_tuning_service.ListTuningJobsResponse: """Post-rpc interceptor for list_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_list_tuning_jobs` interceptor runs + before the `post_list_tuning_jobs_with_metadata` interceptor. """ return response + def post_list_tuning_jobs_with_metadata( + self, + response: genai_tuning_service.ListTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + genai_tuning_service.ListTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_list_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_tuning_jobs_with_metadata` interceptor runs after the + `post_list_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_tuning_jobs` will be passed to + `post_list_tuning_jobs_with_metadata`. + """ + return response, metadata + def pre_rebase_tuned_model( self, request: genai_tuning_service.RebaseTunedModelRequest, @@ -237,12 +309,35 @@ def post_rebase_tuned_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for rebase_tuned_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_rebase_tuned_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_rebase_tuned_model` interceptor runs + before the `post_rebase_tuned_model_with_metadata` interceptor. """ return response + def post_rebase_tuned_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for rebase_tuned_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_rebase_tuned_model_with_metadata` + interceptor in new development instead of the `post_rebase_tuned_model` interceptor. + When both interceptors are used, this `post_rebase_tuned_model_with_metadata` interceptor runs after the + `post_rebase_tuned_model` interceptor. The (possibly modified) response returned by + `post_rebase_tuned_model` will be passed to + `post_rebase_tuned_model_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2907,6 +3002,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3052,6 +3151,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3199,6 +3302,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3353,6 +3460,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_rebase_tuned_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_rebase_tuned_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest_asyncio.py index 07cbe30639..54630a9204 100644 --- a/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/gen_ai_tuning_service/transports/rest_asyncio.py @@ -177,12 +177,35 @@ async def post_create_tuning_job( ) -> gca_tuning_job.TuningJob: """Post-rpc interceptor for create_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_create_tuning_job` interceptor runs + before the `post_create_tuning_job_with_metadata` interceptor. """ return response + async def post_create_tuning_job_with_metadata( + self, + response: gca_tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_create_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_tuning_job` interceptor. + When both interceptors are used, this `post_create_tuning_job_with_metadata` interceptor runs after the + `post_create_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_tuning_job` will be passed to + `post_create_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_get_tuning_job( self, request: genai_tuning_service.GetTuningJobRequest, @@ -203,12 +226,35 @@ async def post_get_tuning_job( ) -> tuning_job.TuningJob: """Post-rpc interceptor for get_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_get_tuning_job` interceptor runs + before the `post_get_tuning_job_with_metadata` interceptor. """ return response + async def post_get_tuning_job_with_metadata( + self, + response: tuning_job.TuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tuning_job.TuningJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_get_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_tuning_job` interceptor. + When both interceptors are used, this `post_get_tuning_job_with_metadata` interceptor runs after the + `post_get_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_tuning_job` will be passed to + `post_get_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_list_tuning_jobs( self, request: genai_tuning_service.ListTuningJobsRequest, @@ -229,12 +275,38 @@ async def post_list_tuning_jobs( ) -> genai_tuning_service.ListTuningJobsResponse: """Post-rpc interceptor for list_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_list_tuning_jobs` interceptor runs + before the `post_list_tuning_jobs_with_metadata` interceptor. """ return response + async def post_list_tuning_jobs_with_metadata( + self, + response: genai_tuning_service.ListTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + genai_tuning_service.ListTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_list_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_tuning_jobs_with_metadata` interceptor runs after the + `post_list_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_tuning_jobs` will be passed to + `post_list_tuning_jobs_with_metadata`. + """ + return response, metadata + async def pre_rebase_tuned_model( self, request: genai_tuning_service.RebaseTunedModelRequest, @@ -255,12 +327,35 @@ async def post_rebase_tuned_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for rebase_tuned_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_rebase_tuned_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the GenAiTuningService server but before - it is returned to user code. + it is returned to user code. This `post_rebase_tuned_model` interceptor runs + before the `post_rebase_tuned_model_with_metadata` interceptor. """ return response + async def post_rebase_tuned_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for rebase_tuned_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the GenAiTuningService server but before it is returned to user code. + + We recommend only using this `post_rebase_tuned_model_with_metadata` + interceptor in new development instead of the `post_rebase_tuned_model` interceptor. + When both interceptors are used, this `post_rebase_tuned_model_with_metadata` interceptor runs after the + `post_rebase_tuned_model` interceptor. The (possibly modified) response returned by + `post_rebase_tuned_model` will be passed to + `post_rebase_tuned_model_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -916,6 +1011,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1071,6 +1170,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1224,6 +1327,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1386,6 +1493,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_rebase_tuned_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_rebase_tuned_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/client.py b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/client.py index 1a9d997b2a..c292b9b615 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -555,6 +557,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1879,16 +1908,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1934,16 +1967,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2106,16 +2143,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2227,16 +2268,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2349,16 +2394,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2409,16 +2458,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2464,16 +2517,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2519,16 +2576,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest.py index 1ec05401b7..1265793b23 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest.py @@ -171,12 +171,35 @@ def post_create_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_index_endpoint` interceptor runs + before the `post_create_index_endpoint_with_metadata` interceptor. """ return response + def post_create_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_create_index_endpoint_with_metadata` + interceptor in new development instead of the `post_create_index_endpoint` interceptor. + When both interceptors are used, this `post_create_index_endpoint_with_metadata` interceptor runs after the + `post_create_index_endpoint` interceptor. The (possibly modified) response returned by + `post_create_index_endpoint` will be passed to + `post_create_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_delete_index_endpoint( self, request: index_endpoint_service.DeleteIndexEndpointRequest, @@ -197,12 +220,35 @@ def post_delete_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index_endpoint` interceptor runs + before the `post_delete_index_endpoint_with_metadata` interceptor. """ return response + def post_delete_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_index_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_index_endpoint` interceptor. + When both interceptors are used, this `post_delete_index_endpoint_with_metadata` interceptor runs after the + `post_delete_index_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_index_endpoint` will be passed to + `post_delete_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_deploy_index( self, request: index_endpoint_service.DeployIndexRequest, @@ -223,12 +269,35 @@ def post_deploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_index` interceptor runs + before the `post_deploy_index_with_metadata` interceptor. """ return response + def post_deploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_index_with_metadata` + interceptor in new development instead of the `post_deploy_index` interceptor. + When both interceptors are used, this `post_deploy_index_with_metadata` interceptor runs after the + `post_deploy_index` interceptor. The (possibly modified) response returned by + `post_deploy_index` will be passed to + `post_deploy_index_with_metadata`. + """ + return response, metadata + def pre_get_index_endpoint( self, request: index_endpoint_service.GetIndexEndpointRequest, @@ -249,12 +318,35 @@ def post_get_index_endpoint( ) -> index_endpoint.IndexEndpoint: """Post-rpc interceptor for get_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_index_endpoint` interceptor runs + before the `post_get_index_endpoint_with_metadata` interceptor. """ return response + def post_get_index_endpoint_with_metadata( + self, + response: index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_get_index_endpoint_with_metadata` + interceptor in new development instead of the `post_get_index_endpoint` interceptor. + When both interceptors are used, this `post_get_index_endpoint_with_metadata` interceptor runs after the + `post_get_index_endpoint` interceptor. The (possibly modified) response returned by + `post_get_index_endpoint` will be passed to + `post_get_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_list_index_endpoints( self, request: index_endpoint_service.ListIndexEndpointsRequest, @@ -275,12 +367,38 @@ def post_list_index_endpoints( ) -> index_endpoint_service.ListIndexEndpointsResponse: """Post-rpc interceptor for list_index_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_index_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_index_endpoints` interceptor runs + before the `post_list_index_endpoints_with_metadata` interceptor. """ return response + def post_list_index_endpoints_with_metadata( + self, + response: index_endpoint_service.ListIndexEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_endpoint_service.ListIndexEndpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_index_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_list_index_endpoints_with_metadata` + interceptor in new development instead of the `post_list_index_endpoints` interceptor. + When both interceptors are used, this `post_list_index_endpoints_with_metadata` interceptor runs after the + `post_list_index_endpoints` interceptor. The (possibly modified) response returned by + `post_list_index_endpoints` will be passed to + `post_list_index_endpoints_with_metadata`. + """ + return response, metadata + def pre_mutate_deployed_index( self, request: index_endpoint_service.MutateDeployedIndexRequest, @@ -301,12 +419,35 @@ def post_mutate_deployed_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_index` interceptor runs + before the `post_mutate_deployed_index_with_metadata` interceptor. """ return response + def post_mutate_deployed_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_index_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_index` interceptor. + When both interceptors are used, this `post_mutate_deployed_index_with_metadata` interceptor runs after the + `post_mutate_deployed_index` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_index` will be passed to + `post_mutate_deployed_index_with_metadata`. + """ + return response, metadata + def pre_undeploy_index( self, request: index_endpoint_service.UndeployIndexRequest, @@ -327,12 +468,35 @@ def post_undeploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_index` interceptor runs + before the `post_undeploy_index_with_metadata` interceptor. """ return response + def post_undeploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_index_with_metadata` + interceptor in new development instead of the `post_undeploy_index` interceptor. + When both interceptors are used, this `post_undeploy_index_with_metadata` interceptor runs after the + `post_undeploy_index` interceptor. The (possibly modified) response returned by + `post_undeploy_index` will be passed to + `post_undeploy_index_with_metadata`. + """ + return response, metadata + def pre_update_index_endpoint( self, request: index_endpoint_service.UpdateIndexEndpointRequest, @@ -353,12 +517,37 @@ def post_update_index_endpoint( ) -> gca_index_endpoint.IndexEndpoint: """Post-rpc interceptor for update_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_index_endpoint` interceptor runs + before the `post_update_index_endpoint_with_metadata` interceptor. """ return response + def post_update_index_endpoint_with_metadata( + self, + response: gca_index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_update_index_endpoint_with_metadata` + interceptor in new development instead of the `post_update_index_endpoint` interceptor. + When both interceptors are used, this `post_update_index_endpoint_with_metadata` interceptor runs after the + `post_update_index_endpoint` interceptor. The (possibly modified) response returned by + `post_update_index_endpoint` will be passed to + `post_update_index_endpoint_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2907,6 +3096,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3055,6 +3248,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3205,6 +3402,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_deploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_deploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3355,6 +3556,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3504,6 +3709,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_index_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_index_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3662,6 +3871,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_mutate_deployed_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_mutate_deployed_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3812,6 +4025,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_undeploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_undeploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3968,6 +4185,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest_asyncio.py index de82b44414..31b7cd33dc 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/index_endpoint_service/transports/rest_asyncio.py @@ -189,12 +189,35 @@ async def post_create_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_create_index_endpoint` interceptor runs + before the `post_create_index_endpoint_with_metadata` interceptor. """ return response + async def post_create_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_create_index_endpoint_with_metadata` + interceptor in new development instead of the `post_create_index_endpoint` interceptor. + When both interceptors are used, this `post_create_index_endpoint_with_metadata` interceptor runs after the + `post_create_index_endpoint` interceptor. The (possibly modified) response returned by + `post_create_index_endpoint` will be passed to + `post_create_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_delete_index_endpoint( self, request: index_endpoint_service.DeleteIndexEndpointRequest, @@ -215,12 +238,35 @@ async def post_delete_index_endpoint( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index_endpoint` interceptor runs + before the `post_delete_index_endpoint_with_metadata` interceptor. """ return response + async def post_delete_index_endpoint_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_delete_index_endpoint_with_metadata` + interceptor in new development instead of the `post_delete_index_endpoint` interceptor. + When both interceptors are used, this `post_delete_index_endpoint_with_metadata` interceptor runs after the + `post_delete_index_endpoint` interceptor. The (possibly modified) response returned by + `post_delete_index_endpoint` will be passed to + `post_delete_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_deploy_index( self, request: index_endpoint_service.DeployIndexRequest, @@ -241,12 +287,35 @@ async def post_deploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_index` interceptor runs + before the `post_deploy_index_with_metadata` interceptor. """ return response + async def post_deploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_deploy_index_with_metadata` + interceptor in new development instead of the `post_deploy_index` interceptor. + When both interceptors are used, this `post_deploy_index_with_metadata` interceptor runs after the + `post_deploy_index` interceptor. The (possibly modified) response returned by + `post_deploy_index` will be passed to + `post_deploy_index_with_metadata`. + """ + return response, metadata + async def pre_get_index_endpoint( self, request: index_endpoint_service.GetIndexEndpointRequest, @@ -267,12 +336,35 @@ async def post_get_index_endpoint( ) -> index_endpoint.IndexEndpoint: """Post-rpc interceptor for get_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_get_index_endpoint` interceptor runs + before the `post_get_index_endpoint_with_metadata` interceptor. """ return response + async def post_get_index_endpoint_with_metadata( + self, + response: index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_get_index_endpoint_with_metadata` + interceptor in new development instead of the `post_get_index_endpoint` interceptor. + When both interceptors are used, this `post_get_index_endpoint_with_metadata` interceptor runs after the + `post_get_index_endpoint` interceptor. The (possibly modified) response returned by + `post_get_index_endpoint` will be passed to + `post_get_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_list_index_endpoints( self, request: index_endpoint_service.ListIndexEndpointsRequest, @@ -293,12 +385,38 @@ async def post_list_index_endpoints( ) -> index_endpoint_service.ListIndexEndpointsResponse: """Post-rpc interceptor for list_index_endpoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_index_endpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_list_index_endpoints` interceptor runs + before the `post_list_index_endpoints_with_metadata` interceptor. """ return response + async def post_list_index_endpoints_with_metadata( + self, + response: index_endpoint_service.ListIndexEndpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_endpoint_service.ListIndexEndpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_index_endpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_list_index_endpoints_with_metadata` + interceptor in new development instead of the `post_list_index_endpoints` interceptor. + When both interceptors are used, this `post_list_index_endpoints_with_metadata` interceptor runs after the + `post_list_index_endpoints` interceptor. The (possibly modified) response returned by + `post_list_index_endpoints` will be passed to + `post_list_index_endpoints_with_metadata`. + """ + return response, metadata + async def pre_mutate_deployed_index( self, request: index_endpoint_service.MutateDeployedIndexRequest, @@ -319,12 +437,35 @@ async def post_mutate_deployed_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for mutate_deployed_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_mutate_deployed_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_mutate_deployed_index` interceptor runs + before the `post_mutate_deployed_index_with_metadata` interceptor. """ return response + async def post_mutate_deployed_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for mutate_deployed_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_mutate_deployed_index_with_metadata` + interceptor in new development instead of the `post_mutate_deployed_index` interceptor. + When both interceptors are used, this `post_mutate_deployed_index_with_metadata` interceptor runs after the + `post_mutate_deployed_index` interceptor. The (possibly modified) response returned by + `post_mutate_deployed_index` will be passed to + `post_mutate_deployed_index_with_metadata`. + """ + return response, metadata + async def pre_undeploy_index( self, request: index_endpoint_service.UndeployIndexRequest, @@ -345,12 +486,35 @@ async def post_undeploy_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for undeploy_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_undeploy_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_undeploy_index` interceptor runs + before the `post_undeploy_index_with_metadata` interceptor. """ return response + async def post_undeploy_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for undeploy_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_undeploy_index_with_metadata` + interceptor in new development instead of the `post_undeploy_index` interceptor. + When both interceptors are used, this `post_undeploy_index_with_metadata` interceptor runs after the + `post_undeploy_index` interceptor. The (possibly modified) response returned by + `post_undeploy_index` will be passed to + `post_undeploy_index_with_metadata`. + """ + return response, metadata + async def pre_update_index_endpoint( self, request: index_endpoint_service.UpdateIndexEndpointRequest, @@ -371,12 +535,37 @@ async def post_update_index_endpoint( ) -> gca_index_endpoint.IndexEndpoint: """Post-rpc interceptor for update_index_endpoint - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_endpoint_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexEndpointService server but before - it is returned to user code. + it is returned to user code. This `post_update_index_endpoint` interceptor runs + before the `post_update_index_endpoint_with_metadata` interceptor. """ return response + async def post_update_index_endpoint_with_metadata( + self, + response: gca_index_endpoint.IndexEndpoint, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_index_endpoint.IndexEndpoint, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_index_endpoint + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexEndpointService server but before it is returned to user code. + + We recommend only using this `post_update_index_endpoint_with_metadata` + interceptor in new development instead of the `post_update_index_endpoint` interceptor. + When both interceptors are used, this `post_update_index_endpoint_with_metadata` interceptor runs after the + `post_update_index_endpoint` interceptor. The (possibly modified) response returned by + `post_update_index_endpoint` will be passed to + `post_update_index_endpoint_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -925,6 +1114,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1079,6 +1272,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1241,6 +1438,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_deploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_deploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1395,6 +1596,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1548,6 +1753,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_index_endpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_index_endpoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1712,6 +1921,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_mutate_deployed_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_mutate_deployed_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1872,6 +2085,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_undeploy_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_undeploy_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2032,6 +2249,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_index_endpoint(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_index_endpoint_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/index_service/client.py b/google/cloud/aiplatform_v1beta1/services/index_service/client.py index 0865754eb3..c438b4fe9a 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/index_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -534,6 +536,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1604,16 +1633,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1659,16 +1692,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1831,16 +1868,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1952,16 +1993,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2074,16 +2119,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2134,16 +2183,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2189,16 +2242,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2244,16 +2301,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest.py index a612bffb10..905f211e67 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest.py @@ -161,12 +161,35 @@ def post_create_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_create_index` interceptor runs + before the `post_create_index_with_metadata` interceptor. """ return response + def post_create_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_create_index_with_metadata` + interceptor in new development instead of the `post_create_index` interceptor. + When both interceptors are used, this `post_create_index_with_metadata` interceptor runs after the + `post_create_index` interceptor. The (possibly modified) response returned by + `post_create_index` will be passed to + `post_create_index_with_metadata`. + """ + return response, metadata + def pre_delete_index( self, request: index_service.DeleteIndexRequest, @@ -186,12 +209,35 @@ def post_delete_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index` interceptor runs + before the `post_delete_index_with_metadata` interceptor. """ return response + def post_delete_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_delete_index_with_metadata` + interceptor in new development instead of the `post_delete_index` interceptor. + When both interceptors are used, this `post_delete_index_with_metadata` interceptor runs after the + `post_delete_index` interceptor. The (possibly modified) response returned by + `post_delete_index` will be passed to + `post_delete_index_with_metadata`. + """ + return response, metadata + def pre_get_index( self, request: index_service.GetIndexRequest, @@ -207,12 +253,33 @@ def pre_get_index( def post_get_index(self, response: index.Index) -> index.Index: """Post-rpc interceptor for get_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_get_index` interceptor runs + before the `post_get_index_with_metadata` interceptor. """ return response + def post_get_index_with_metadata( + self, response: index.Index, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[index.Index, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_get_index_with_metadata` + interceptor in new development instead of the `post_get_index` interceptor. + When both interceptors are used, this `post_get_index_with_metadata` interceptor runs after the + `post_get_index` interceptor. The (possibly modified) response returned by + `post_get_index` will be passed to + `post_get_index_with_metadata`. + """ + return response, metadata + def pre_list_indexes( self, request: index_service.ListIndexesRequest, @@ -232,12 +299,37 @@ def post_list_indexes( ) -> index_service.ListIndexesResponse: """Post-rpc interceptor for list_indexes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_indexes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_list_indexes` interceptor runs + before the `post_list_indexes_with_metadata` interceptor. """ return response + def post_list_indexes_with_metadata( + self, + response: index_service.ListIndexesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.ListIndexesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_indexes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_list_indexes_with_metadata` + interceptor in new development instead of the `post_list_indexes` interceptor. + When both interceptors are used, this `post_list_indexes_with_metadata` interceptor runs after the + `post_list_indexes` interceptor. The (possibly modified) response returned by + `post_list_indexes` will be passed to + `post_list_indexes_with_metadata`. + """ + return response, metadata + def pre_remove_datapoints( self, request: index_service.RemoveDatapointsRequest, @@ -257,12 +349,37 @@ def post_remove_datapoints( ) -> index_service.RemoveDatapointsResponse: """Post-rpc interceptor for remove_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_remove_datapoints` interceptor runs + before the `post_remove_datapoints_with_metadata` interceptor. """ return response + def post_remove_datapoints_with_metadata( + self, + response: index_service.RemoveDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.RemoveDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for remove_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_remove_datapoints_with_metadata` + interceptor in new development instead of the `post_remove_datapoints` interceptor. + When both interceptors are used, this `post_remove_datapoints_with_metadata` interceptor runs after the + `post_remove_datapoints` interceptor. The (possibly modified) response returned by + `post_remove_datapoints` will be passed to + `post_remove_datapoints_with_metadata`. + """ + return response, metadata + def pre_update_index( self, request: index_service.UpdateIndexRequest, @@ -282,12 +399,35 @@ def post_update_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_update_index` interceptor runs + before the `post_update_index_with_metadata` interceptor. """ return response + def post_update_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_update_index_with_metadata` + interceptor in new development instead of the `post_update_index` interceptor. + When both interceptors are used, this `post_update_index_with_metadata` interceptor runs after the + `post_update_index` interceptor. The (possibly modified) response returned by + `post_update_index` will be passed to + `post_update_index_with_metadata`. + """ + return response, metadata + def pre_upsert_datapoints( self, request: index_service.UpsertDatapointsRequest, @@ -307,12 +447,37 @@ def post_upsert_datapoints( ) -> index_service.UpsertDatapointsResponse: """Post-rpc interceptor for upsert_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upsert_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_upsert_datapoints` interceptor runs + before the `post_upsert_datapoints_with_metadata` interceptor. """ return response + def post_upsert_datapoints_with_metadata( + self, + response: index_service.UpsertDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.UpsertDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for upsert_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_upsert_datapoints_with_metadata` + interceptor in new development instead of the `post_upsert_datapoints` interceptor. + When both interceptors are used, this `post_upsert_datapoints_with_metadata` interceptor runs after the + `post_upsert_datapoints` interceptor. The (possibly modified) response returned by + `post_upsert_datapoints` will be passed to + `post_upsert_datapoints_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2863,6 +3028,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3010,6 +3179,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3158,6 +3331,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3306,6 +3483,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_indexes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_indexes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3460,6 +3641,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_remove_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_remove_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3617,6 +3802,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3769,6 +3958,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_upsert_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upsert_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest_asyncio.py index b8d4a2324a..054cd36c2b 100644 --- a/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/index_service/transports/rest_asyncio.py @@ -179,12 +179,35 @@ async def post_create_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_create_index` interceptor runs + before the `post_create_index_with_metadata` interceptor. """ return response + async def post_create_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_create_index_with_metadata` + interceptor in new development instead of the `post_create_index` interceptor. + When both interceptors are used, this `post_create_index_with_metadata` interceptor runs after the + `post_create_index` interceptor. The (possibly modified) response returned by + `post_create_index` will be passed to + `post_create_index_with_metadata`. + """ + return response, metadata + async def pre_delete_index( self, request: index_service.DeleteIndexRequest, @@ -204,12 +227,35 @@ async def post_delete_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_delete_index` interceptor runs + before the `post_delete_index_with_metadata` interceptor. """ return response + async def post_delete_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_delete_index_with_metadata` + interceptor in new development instead of the `post_delete_index` interceptor. + When both interceptors are used, this `post_delete_index_with_metadata` interceptor runs after the + `post_delete_index` interceptor. The (possibly modified) response returned by + `post_delete_index` will be passed to + `post_delete_index_with_metadata`. + """ + return response, metadata + async def pre_get_index( self, request: index_service.GetIndexRequest, @@ -225,12 +271,33 @@ async def pre_get_index( async def post_get_index(self, response: index.Index) -> index.Index: """Post-rpc interceptor for get_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_get_index` interceptor runs + before the `post_get_index_with_metadata` interceptor. """ return response + async def post_get_index_with_metadata( + self, response: index.Index, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[index.Index, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_get_index_with_metadata` + interceptor in new development instead of the `post_get_index` interceptor. + When both interceptors are used, this `post_get_index_with_metadata` interceptor runs after the + `post_get_index` interceptor. The (possibly modified) response returned by + `post_get_index` will be passed to + `post_get_index_with_metadata`. + """ + return response, metadata + async def pre_list_indexes( self, request: index_service.ListIndexesRequest, @@ -250,12 +317,37 @@ async def post_list_indexes( ) -> index_service.ListIndexesResponse: """Post-rpc interceptor for list_indexes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_indexes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_list_indexes` interceptor runs + before the `post_list_indexes_with_metadata` interceptor. """ return response + async def post_list_indexes_with_metadata( + self, + response: index_service.ListIndexesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.ListIndexesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_indexes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_list_indexes_with_metadata` + interceptor in new development instead of the `post_list_indexes` interceptor. + When both interceptors are used, this `post_list_indexes_with_metadata` interceptor runs after the + `post_list_indexes` interceptor. The (possibly modified) response returned by + `post_list_indexes` will be passed to + `post_list_indexes_with_metadata`. + """ + return response, metadata + async def pre_remove_datapoints( self, request: index_service.RemoveDatapointsRequest, @@ -275,12 +367,37 @@ async def post_remove_datapoints( ) -> index_service.RemoveDatapointsResponse: """Post-rpc interceptor for remove_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_remove_datapoints` interceptor runs + before the `post_remove_datapoints_with_metadata` interceptor. """ return response + async def post_remove_datapoints_with_metadata( + self, + response: index_service.RemoveDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.RemoveDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for remove_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_remove_datapoints_with_metadata` + interceptor in new development instead of the `post_remove_datapoints` interceptor. + When both interceptors are used, this `post_remove_datapoints_with_metadata` interceptor runs after the + `post_remove_datapoints` interceptor. The (possibly modified) response returned by + `post_remove_datapoints` will be passed to + `post_remove_datapoints_with_metadata`. + """ + return response, metadata + async def pre_update_index( self, request: index_service.UpdateIndexRequest, @@ -300,12 +417,35 @@ async def post_update_index( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_index - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_index_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_update_index` interceptor runs + before the `post_update_index_with_metadata` interceptor. """ return response + async def post_update_index_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_index + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_update_index_with_metadata` + interceptor in new development instead of the `post_update_index` interceptor. + When both interceptors are used, this `post_update_index_with_metadata` interceptor runs after the + `post_update_index` interceptor. The (possibly modified) response returned by + `post_update_index` will be passed to + `post_update_index_with_metadata`. + """ + return response, metadata + async def pre_upsert_datapoints( self, request: index_service.UpsertDatapointsRequest, @@ -325,12 +465,37 @@ async def post_upsert_datapoints( ) -> index_service.UpsertDatapointsResponse: """Post-rpc interceptor for upsert_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upsert_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the IndexService server but before - it is returned to user code. + it is returned to user code. This `post_upsert_datapoints` interceptor runs + before the `post_upsert_datapoints_with_metadata` interceptor. """ return response + async def post_upsert_datapoints_with_metadata( + self, + response: index_service.UpsertDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + index_service.UpsertDatapointsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for upsert_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the IndexService server but before it is returned to user code. + + We recommend only using this `post_upsert_datapoints_with_metadata` + interceptor in new development instead of the `post_upsert_datapoints` interceptor. + When both interceptors are used, this `post_upsert_datapoints_with_metadata` interceptor runs after the + `post_upsert_datapoints` interceptor. The (possibly modified) response returned by + `post_upsert_datapoints` will be passed to + `post_upsert_datapoints_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -880,6 +1045,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1037,6 +1206,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1193,6 +1366,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1349,6 +1526,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_indexes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_indexes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1511,6 +1692,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_remove_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_remove_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1678,6 +1863,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_index(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_index_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1838,6 +2027,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upsert_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upsert_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/job_service/client.py b/google/cloud/aiplatform_v1beta1/services/job_service/client.py index db1cfa15e3..f8ee955239 100644 --- a/google/cloud/aiplatform_v1beta1/services/job_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/job_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -916,6 +918,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -5437,16 +5466,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -5492,16 +5525,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -5664,16 +5701,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5785,16 +5826,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5907,16 +5952,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5967,16 +6016,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -6022,16 +6075,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -6077,16 +6134,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest.py index 422e2aaf01..b91d8c0f65 100644 --- a/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest.py @@ -451,12 +451,38 @@ def post_create_batch_prediction_job( ) -> gca_batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for create_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_batch_prediction_job` interceptor runs + before the `post_create_batch_prediction_job_with_metadata` interceptor. """ return response + def post_create_batch_prediction_job_with_metadata( + self, + response: gca_batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_batch_prediction_job.BatchPredictionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_create_batch_prediction_job` interceptor. + When both interceptors are used, this `post_create_batch_prediction_job_with_metadata` interceptor runs after the + `post_create_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_create_batch_prediction_job` will be passed to + `post_create_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_create_custom_job( self, request: job_service.CreateCustomJobRequest, @@ -476,12 +502,35 @@ def post_create_custom_job( ) -> gca_custom_job.CustomJob: """Post-rpc interceptor for create_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_custom_job` interceptor runs + before the `post_create_custom_job_with_metadata` interceptor. """ return response + def post_create_custom_job_with_metadata( + self, + response: gca_custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_custom_job_with_metadata` + interceptor in new development instead of the `post_create_custom_job` interceptor. + When both interceptors are used, this `post_create_custom_job_with_metadata` interceptor runs after the + `post_create_custom_job` interceptor. The (possibly modified) response returned by + `post_create_custom_job` will be passed to + `post_create_custom_job_with_metadata`. + """ + return response, metadata + def pre_create_data_labeling_job( self, request: job_service.CreateDataLabelingJobRequest, @@ -502,12 +551,37 @@ def post_create_data_labeling_job( ) -> gca_data_labeling_job.DataLabelingJob: """Post-rpc interceptor for create_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_data_labeling_job` interceptor runs + before the `post_create_data_labeling_job_with_metadata` interceptor. """ return response + def post_create_data_labeling_job_with_metadata( + self, + response: gca_data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_create_data_labeling_job` interceptor. + When both interceptors are used, this `post_create_data_labeling_job_with_metadata` interceptor runs after the + `post_create_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_create_data_labeling_job` will be passed to + `post_create_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_create_hyperparameter_tuning_job( self, request: job_service.CreateHyperparameterTuningJobRequest, @@ -528,12 +602,38 @@ def post_create_hyperparameter_tuning_job( ) -> gca_hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for create_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_hyperparameter_tuning_job` interceptor runs + before the `post_create_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_create_hyperparameter_tuning_job_with_metadata( + self, + response: gca_hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_create_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_create_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_hyperparameter_tuning_job` will be passed to + `post_create_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_create_model_deployment_monitoring_job( self, request: job_service.CreateModelDeploymentMonitoringJobRequest, @@ -554,12 +654,38 @@ def post_create_model_deployment_monitoring_job( ) -> gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for create_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_deployment_monitoring_job` interceptor runs + before the `post_create_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_create_model_deployment_monitoring_job_with_metadata( + self, + response: gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_deployment_monitoring_job` will be passed to + `post_create_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_create_nas_job( self, request: job_service.CreateNasJobRequest, @@ -577,12 +703,35 @@ def pre_create_nas_job( def post_create_nas_job(self, response: gca_nas_job.NasJob) -> gca_nas_job.NasJob: """Post-rpc interceptor for create_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_nas_job` interceptor runs + before the `post_create_nas_job_with_metadata` interceptor. """ return response + def post_create_nas_job_with_metadata( + self, + response: gca_nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_nas_job_with_metadata` + interceptor in new development instead of the `post_create_nas_job` interceptor. + When both interceptors are used, this `post_create_nas_job_with_metadata` interceptor runs after the + `post_create_nas_job` interceptor. The (possibly modified) response returned by + `post_create_nas_job` will be passed to + `post_create_nas_job_with_metadata`. + """ + return response, metadata + def pre_delete_batch_prediction_job( self, request: job_service.DeleteBatchPredictionJobRequest, @@ -603,12 +752,35 @@ def post_delete_batch_prediction_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_batch_prediction_job` interceptor runs + before the `post_delete_batch_prediction_job_with_metadata` interceptor. """ return response + def post_delete_batch_prediction_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_delete_batch_prediction_job` interceptor. + When both interceptors are used, this `post_delete_batch_prediction_job_with_metadata` interceptor runs after the + `post_delete_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_delete_batch_prediction_job` will be passed to + `post_delete_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_delete_custom_job( self, request: job_service.DeleteCustomJobRequest, @@ -628,12 +800,35 @@ def post_delete_custom_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_custom_job` interceptor runs + before the `post_delete_custom_job_with_metadata` interceptor. """ return response + def post_delete_custom_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_custom_job_with_metadata` + interceptor in new development instead of the `post_delete_custom_job` interceptor. + When both interceptors are used, this `post_delete_custom_job_with_metadata` interceptor runs after the + `post_delete_custom_job` interceptor. The (possibly modified) response returned by + `post_delete_custom_job` will be passed to + `post_delete_custom_job_with_metadata`. + """ + return response, metadata + def pre_delete_data_labeling_job( self, request: job_service.DeleteDataLabelingJobRequest, @@ -654,12 +849,35 @@ def post_delete_data_labeling_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_data_labeling_job` interceptor runs + before the `post_delete_data_labeling_job_with_metadata` interceptor. """ return response + def post_delete_data_labeling_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_delete_data_labeling_job` interceptor. + When both interceptors are used, this `post_delete_data_labeling_job_with_metadata` interceptor runs after the + `post_delete_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_delete_data_labeling_job` will be passed to + `post_delete_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_delete_hyperparameter_tuning_job( self, request: job_service.DeleteHyperparameterTuningJobRequest, @@ -680,12 +898,35 @@ def post_delete_hyperparameter_tuning_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_hyperparameter_tuning_job` interceptor runs + before the `post_delete_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_delete_hyperparameter_tuning_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_delete_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_delete_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_delete_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_delete_hyperparameter_tuning_job` will be passed to + `post_delete_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_delete_model_deployment_monitoring_job( self, request: job_service.DeleteModelDeploymentMonitoringJobRequest, @@ -706,12 +947,35 @@ def post_delete_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_deployment_monitoring_job` interceptor runs + before the `post_delete_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_delete_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_deployment_monitoring_job` will be passed to + `post_delete_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_delete_nas_job( self, request: job_service.DeleteNasJobRequest, @@ -731,12 +995,35 @@ def post_delete_nas_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_nas_job` interceptor runs + before the `post_delete_nas_job_with_metadata` interceptor. """ return response + def post_delete_nas_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_nas_job_with_metadata` + interceptor in new development instead of the `post_delete_nas_job` interceptor. + When both interceptors are used, this `post_delete_nas_job_with_metadata` interceptor runs after the + `post_delete_nas_job` interceptor. The (possibly modified) response returned by + `post_delete_nas_job` will be passed to + `post_delete_nas_job_with_metadata`. + """ + return response, metadata + def pre_get_batch_prediction_job( self, request: job_service.GetBatchPredictionJobRequest, @@ -757,12 +1044,37 @@ def post_get_batch_prediction_job( ) -> batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for get_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_batch_prediction_job` interceptor runs + before the `post_get_batch_prediction_job_with_metadata` interceptor. """ return response + def post_get_batch_prediction_job_with_metadata( + self, + response: batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + batch_prediction_job.BatchPredictionJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_get_batch_prediction_job` interceptor. + When both interceptors are used, this `post_get_batch_prediction_job_with_metadata` interceptor runs after the + `post_get_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_get_batch_prediction_job` will be passed to + `post_get_batch_prediction_job_with_metadata`. + """ + return response, metadata + def pre_get_custom_job( self, request: job_service.GetCustomJobRequest, @@ -782,12 +1094,35 @@ def post_get_custom_job( ) -> custom_job.CustomJob: """Post-rpc interceptor for get_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_custom_job` interceptor runs + before the `post_get_custom_job_with_metadata` interceptor. """ return response + def post_get_custom_job_with_metadata( + self, + response: custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_custom_job_with_metadata` + interceptor in new development instead of the `post_get_custom_job` interceptor. + When both interceptors are used, this `post_get_custom_job_with_metadata` interceptor runs after the + `post_get_custom_job` interceptor. The (possibly modified) response returned by + `post_get_custom_job` will be passed to + `post_get_custom_job_with_metadata`. + """ + return response, metadata + def pre_get_data_labeling_job( self, request: job_service.GetDataLabelingJobRequest, @@ -807,12 +1142,37 @@ def post_get_data_labeling_job( ) -> data_labeling_job.DataLabelingJob: """Post-rpc interceptor for get_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_data_labeling_job` interceptor runs + before the `post_get_data_labeling_job_with_metadata` interceptor. """ return response + def post_get_data_labeling_job_with_metadata( + self, + response: data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_get_data_labeling_job` interceptor. + When both interceptors are used, this `post_get_data_labeling_job_with_metadata` interceptor runs after the + `post_get_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_get_data_labeling_job` will be passed to + `post_get_data_labeling_job_with_metadata`. + """ + return response, metadata + def pre_get_hyperparameter_tuning_job( self, request: job_service.GetHyperparameterTuningJobRequest, @@ -833,12 +1193,38 @@ def post_get_hyperparameter_tuning_job( ) -> hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for get_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_hyperparameter_tuning_job` interceptor runs + before the `post_get_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + def post_get_hyperparameter_tuning_job_with_metadata( + self, + response: hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_get_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_get_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_hyperparameter_tuning_job` will be passed to + `post_get_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + def pre_get_model_deployment_monitoring_job( self, request: job_service.GetModelDeploymentMonitoringJobRequest, @@ -859,12 +1245,38 @@ def post_get_model_deployment_monitoring_job( ) -> model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for get_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_deployment_monitoring_job` interceptor runs + before the `post_get_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_get_model_deployment_monitoring_job_with_metadata( + self, + response: model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_deployment_monitoring_job` will be passed to + `post_get_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_get_nas_job( self, request: job_service.GetNasJobRequest, @@ -880,12 +1292,35 @@ def pre_get_nas_job( def post_get_nas_job(self, response: nas_job.NasJob) -> nas_job.NasJob: """Post-rpc interceptor for get_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_job` interceptor runs + before the `post_get_nas_job_with_metadata` interceptor. """ return response + def post_get_nas_job_with_metadata( + self, + response: nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_job_with_metadata` + interceptor in new development instead of the `post_get_nas_job` interceptor. + When both interceptors are used, this `post_get_nas_job_with_metadata` interceptor runs after the + `post_get_nas_job` interceptor. The (possibly modified) response returned by + `post_get_nas_job` will be passed to + `post_get_nas_job_with_metadata`. + """ + return response, metadata + def pre_get_nas_trial_detail( self, request: job_service.GetNasTrialDetailRequest, @@ -905,12 +1340,35 @@ def post_get_nas_trial_detail( ) -> nas_job.NasTrialDetail: """Post-rpc interceptor for get_nas_trial_detail - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_trial_detail_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_trial_detail` interceptor runs + before the `post_get_nas_trial_detail_with_metadata` interceptor. """ return response + def post_get_nas_trial_detail_with_metadata( + self, + response: nas_job.NasTrialDetail, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasTrialDetail, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_trial_detail + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_trial_detail_with_metadata` + interceptor in new development instead of the `post_get_nas_trial_detail` interceptor. + When both interceptors are used, this `post_get_nas_trial_detail_with_metadata` interceptor runs after the + `post_get_nas_trial_detail` interceptor. The (possibly modified) response returned by + `post_get_nas_trial_detail` will be passed to + `post_get_nas_trial_detail_with_metadata`. + """ + return response, metadata + def pre_list_batch_prediction_jobs( self, request: job_service.ListBatchPredictionJobsRequest, @@ -931,12 +1389,38 @@ def post_list_batch_prediction_jobs( ) -> job_service.ListBatchPredictionJobsResponse: """Post-rpc interceptor for list_batch_prediction_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_batch_prediction_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_batch_prediction_jobs` interceptor runs + before the `post_list_batch_prediction_jobs_with_metadata` interceptor. """ return response + def post_list_batch_prediction_jobs_with_metadata( + self, + response: job_service.ListBatchPredictionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListBatchPredictionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_batch_prediction_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_batch_prediction_jobs_with_metadata` + interceptor in new development instead of the `post_list_batch_prediction_jobs` interceptor. + When both interceptors are used, this `post_list_batch_prediction_jobs_with_metadata` interceptor runs after the + `post_list_batch_prediction_jobs` interceptor. The (possibly modified) response returned by + `post_list_batch_prediction_jobs` will be passed to + `post_list_batch_prediction_jobs_with_metadata`. + """ + return response, metadata + def pre_list_custom_jobs( self, request: job_service.ListCustomJobsRequest, @@ -956,12 +1440,37 @@ def post_list_custom_jobs( ) -> job_service.ListCustomJobsResponse: """Post-rpc interceptor for list_custom_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_custom_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_custom_jobs` interceptor runs + before the `post_list_custom_jobs_with_metadata` interceptor. """ return response + def post_list_custom_jobs_with_metadata( + self, + response: job_service.ListCustomJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListCustomJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_custom_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_custom_jobs_with_metadata` + interceptor in new development instead of the `post_list_custom_jobs` interceptor. + When both interceptors are used, this `post_list_custom_jobs_with_metadata` interceptor runs after the + `post_list_custom_jobs` interceptor. The (possibly modified) response returned by + `post_list_custom_jobs` will be passed to + `post_list_custom_jobs_with_metadata`. + """ + return response, metadata + def pre_list_data_labeling_jobs( self, request: job_service.ListDataLabelingJobsRequest, @@ -981,12 +1490,38 @@ def post_list_data_labeling_jobs( ) -> job_service.ListDataLabelingJobsResponse: """Post-rpc interceptor for list_data_labeling_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_labeling_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_labeling_jobs` interceptor runs + before the `post_list_data_labeling_jobs_with_metadata` interceptor. """ return response + def post_list_data_labeling_jobs_with_metadata( + self, + response: job_service.ListDataLabelingJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListDataLabelingJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_data_labeling_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_data_labeling_jobs_with_metadata` + interceptor in new development instead of the `post_list_data_labeling_jobs` interceptor. + When both interceptors are used, this `post_list_data_labeling_jobs_with_metadata` interceptor runs after the + `post_list_data_labeling_jobs` interceptor. The (possibly modified) response returned by + `post_list_data_labeling_jobs` will be passed to + `post_list_data_labeling_jobs_with_metadata`. + """ + return response, metadata + def pre_list_hyperparameter_tuning_jobs( self, request: job_service.ListHyperparameterTuningJobsRequest, @@ -1007,12 +1542,38 @@ def post_list_hyperparameter_tuning_jobs( ) -> job_service.ListHyperparameterTuningJobsResponse: """Post-rpc interceptor for list_hyperparameter_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_hyperparameter_tuning_jobs` interceptor runs + before the `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor. """ return response + def post_list_hyperparameter_tuning_jobs_with_metadata( + self, + response: job_service.ListHyperparameterTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListHyperparameterTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_hyperparameter_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_hyperparameter_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor runs after the + `post_list_hyperparameter_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_hyperparameter_tuning_jobs` will be passed to + `post_list_hyperparameter_tuning_jobs_with_metadata`. + """ + return response, metadata + def pre_list_model_deployment_monitoring_jobs( self, request: job_service.ListModelDeploymentMonitoringJobsRequest, @@ -1033,12 +1594,38 @@ def post_list_model_deployment_monitoring_jobs( ) -> job_service.ListModelDeploymentMonitoringJobsResponse: """Post-rpc interceptor for list_model_deployment_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_deployment_monitoring_jobs` interceptor runs + before the `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor. """ return response + def post_list_model_deployment_monitoring_jobs_with_metadata( + self, + response: job_service.ListModelDeploymentMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListModelDeploymentMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_deployment_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_deployment_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_deployment_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_deployment_monitoring_jobs` will be passed to + `post_list_model_deployment_monitoring_jobs_with_metadata`. + """ + return response, metadata + def pre_list_nas_jobs( self, request: job_service.ListNasJobsRequest, @@ -1056,12 +1643,37 @@ def post_list_nas_jobs( ) -> job_service.ListNasJobsResponse: """Post-rpc interceptor for list_nas_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_jobs` interceptor runs + before the `post_list_nas_jobs_with_metadata` interceptor. """ return response + def post_list_nas_jobs_with_metadata( + self, + response: job_service.ListNasJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_jobs_with_metadata` + interceptor in new development instead of the `post_list_nas_jobs` interceptor. + When both interceptors are used, this `post_list_nas_jobs_with_metadata` interceptor runs after the + `post_list_nas_jobs` interceptor. The (possibly modified) response returned by + `post_list_nas_jobs` will be passed to + `post_list_nas_jobs_with_metadata`. + """ + return response, metadata + def pre_list_nas_trial_details( self, request: job_service.ListNasTrialDetailsRequest, @@ -1081,12 +1693,37 @@ def post_list_nas_trial_details( ) -> job_service.ListNasTrialDetailsResponse: """Post-rpc interceptor for list_nas_trial_details - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_trial_details_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_trial_details` interceptor runs + before the `post_list_nas_trial_details_with_metadata` interceptor. """ return response + def post_list_nas_trial_details_with_metadata( + self, + response: job_service.ListNasTrialDetailsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasTrialDetailsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_trial_details + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_trial_details_with_metadata` + interceptor in new development instead of the `post_list_nas_trial_details` interceptor. + When both interceptors are used, this `post_list_nas_trial_details_with_metadata` interceptor runs after the + `post_list_nas_trial_details` interceptor. The (possibly modified) response returned by + `post_list_nas_trial_details` will be passed to + `post_list_nas_trial_details_with_metadata`. + """ + return response, metadata + def pre_pause_model_deployment_monitoring_job( self, request: job_service.PauseModelDeploymentMonitoringJobRequest, @@ -1138,12 +1775,38 @@ def post_search_model_deployment_monitoring_stats_anomalies( ) -> job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse: """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_deployment_monitoring_stats_anomalies` interceptor runs + before the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor. """ return response + def post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + self, + response: job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor in new development instead of the `post_search_model_deployment_monitoring_stats_anomalies` interceptor. + When both interceptors are used, this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor runs after the + `post_search_model_deployment_monitoring_stats_anomalies` interceptor. The (possibly modified) response returned by + `post_search_model_deployment_monitoring_stats_anomalies` will be passed to + `post_search_model_deployment_monitoring_stats_anomalies_with_metadata`. + """ + return response, metadata + def pre_update_model_deployment_monitoring_job( self, request: job_service.UpdateModelDeploymentMonitoringJobRequest, @@ -1164,12 +1827,35 @@ def post_update_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_deployment_monitoring_job` interceptor runs + before the `post_update_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + def post_update_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_update_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_update_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_update_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_update_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_update_model_deployment_monitoring_job` will be passed to + `post_update_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -4311,6 +4997,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4471,6 +5161,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4624,6 +5318,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4784,6 +5482,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4952,6 +5657,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5110,6 +5822,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5256,6 +5972,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5401,6 +6121,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5546,6 +6270,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5695,6 +6423,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5845,6 +6580,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5992,6 +6734,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6143,6 +6889,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6299,6 +7049,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6446,6 +7200,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6599,6 +7357,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6761,6 +7526,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6909,6 +7681,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7057,6 +7833,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_nas_trial_detail(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_nas_trial_detail_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7204,6 +7984,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_batch_prediction_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_batch_prediction_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7354,6 +8138,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_custom_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_custom_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7502,6 +8290,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_data_labeling_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_data_labeling_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7654,6 +8446,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_hyperparameter_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_hyperparameter_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7809,6 +8608,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_deployment_monitoring_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_model_deployment_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7961,6 +8767,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_nas_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_nas_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8107,6 +8917,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_nas_trial_details(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_nas_trial_details_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8516,6 +9330,13 @@ def __call__( resp = self._interceptor.post_search_model_deployment_monitoring_stats_anomalies( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8674,6 +9495,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_model_deployment_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest_asyncio.py index 4a9e488c91..12cb5321a1 100644 --- a/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/job_service/transports/rest_asyncio.py @@ -469,12 +469,38 @@ async def post_create_batch_prediction_job( ) -> gca_batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for create_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_batch_prediction_job` interceptor runs + before the `post_create_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_create_batch_prediction_job_with_metadata( + self, + response: gca_batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_batch_prediction_job.BatchPredictionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_create_batch_prediction_job` interceptor. + When both interceptors are used, this `post_create_batch_prediction_job_with_metadata` interceptor runs after the + `post_create_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_create_batch_prediction_job` will be passed to + `post_create_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_create_custom_job( self, request: job_service.CreateCustomJobRequest, @@ -494,12 +520,35 @@ async def post_create_custom_job( ) -> gca_custom_job.CustomJob: """Post-rpc interceptor for create_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_custom_job` interceptor runs + before the `post_create_custom_job_with_metadata` interceptor. """ return response + async def post_create_custom_job_with_metadata( + self, + response: gca_custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_custom_job_with_metadata` + interceptor in new development instead of the `post_create_custom_job` interceptor. + When both interceptors are used, this `post_create_custom_job_with_metadata` interceptor runs after the + `post_create_custom_job` interceptor. The (possibly modified) response returned by + `post_create_custom_job` will be passed to + `post_create_custom_job_with_metadata`. + """ + return response, metadata + async def pre_create_data_labeling_job( self, request: job_service.CreateDataLabelingJobRequest, @@ -520,12 +569,37 @@ async def post_create_data_labeling_job( ) -> gca_data_labeling_job.DataLabelingJob: """Post-rpc interceptor for create_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_data_labeling_job` interceptor runs + before the `post_create_data_labeling_job_with_metadata` interceptor. """ return response + async def post_create_data_labeling_job_with_metadata( + self, + response: gca_data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_create_data_labeling_job` interceptor. + When both interceptors are used, this `post_create_data_labeling_job_with_metadata` interceptor runs after the + `post_create_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_create_data_labeling_job` will be passed to + `post_create_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_create_hyperparameter_tuning_job( self, request: job_service.CreateHyperparameterTuningJobRequest, @@ -546,12 +620,38 @@ async def post_create_hyperparameter_tuning_job( ) -> gca_hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for create_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_hyperparameter_tuning_job` interceptor runs + before the `post_create_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_create_hyperparameter_tuning_job_with_metadata( + self, + response: gca_hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_create_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_create_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_create_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_create_hyperparameter_tuning_job` will be passed to + `post_create_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_create_model_deployment_monitoring_job( self, request: job_service.CreateModelDeploymentMonitoringJobRequest, @@ -572,12 +672,38 @@ async def post_create_model_deployment_monitoring_job( ) -> gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for create_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_deployment_monitoring_job` interceptor runs + before the `post_create_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_create_model_deployment_monitoring_job_with_metadata( + self, + response: gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_deployment_monitoring_job` will be passed to + `post_create_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_create_nas_job( self, request: job_service.CreateNasJobRequest, @@ -597,12 +723,35 @@ async def post_create_nas_job( ) -> gca_nas_job.NasJob: """Post-rpc interceptor for create_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_create_nas_job` interceptor runs + before the `post_create_nas_job_with_metadata` interceptor. """ return response + async def post_create_nas_job_with_metadata( + self, + response: gca_nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_create_nas_job_with_metadata` + interceptor in new development instead of the `post_create_nas_job` interceptor. + When both interceptors are used, this `post_create_nas_job_with_metadata` interceptor runs after the + `post_create_nas_job` interceptor. The (possibly modified) response returned by + `post_create_nas_job` will be passed to + `post_create_nas_job_with_metadata`. + """ + return response, metadata + async def pre_delete_batch_prediction_job( self, request: job_service.DeleteBatchPredictionJobRequest, @@ -623,12 +772,35 @@ async def post_delete_batch_prediction_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_batch_prediction_job` interceptor runs + before the `post_delete_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_delete_batch_prediction_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_delete_batch_prediction_job` interceptor. + When both interceptors are used, this `post_delete_batch_prediction_job_with_metadata` interceptor runs after the + `post_delete_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_delete_batch_prediction_job` will be passed to + `post_delete_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_delete_custom_job( self, request: job_service.DeleteCustomJobRequest, @@ -648,12 +820,35 @@ async def post_delete_custom_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_custom_job` interceptor runs + before the `post_delete_custom_job_with_metadata` interceptor. """ return response + async def post_delete_custom_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_custom_job_with_metadata` + interceptor in new development instead of the `post_delete_custom_job` interceptor. + When both interceptors are used, this `post_delete_custom_job_with_metadata` interceptor runs after the + `post_delete_custom_job` interceptor. The (possibly modified) response returned by + `post_delete_custom_job` will be passed to + `post_delete_custom_job_with_metadata`. + """ + return response, metadata + async def pre_delete_data_labeling_job( self, request: job_service.DeleteDataLabelingJobRequest, @@ -674,12 +869,35 @@ async def post_delete_data_labeling_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_data_labeling_job` interceptor runs + before the `post_delete_data_labeling_job_with_metadata` interceptor. """ return response + async def post_delete_data_labeling_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_delete_data_labeling_job` interceptor. + When both interceptors are used, this `post_delete_data_labeling_job_with_metadata` interceptor runs after the + `post_delete_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_delete_data_labeling_job` will be passed to + `post_delete_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_delete_hyperparameter_tuning_job( self, request: job_service.DeleteHyperparameterTuningJobRequest, @@ -700,12 +918,35 @@ async def post_delete_hyperparameter_tuning_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_hyperparameter_tuning_job` interceptor runs + before the `post_delete_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_delete_hyperparameter_tuning_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_delete_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_delete_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_delete_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_delete_hyperparameter_tuning_job` will be passed to + `post_delete_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_delete_model_deployment_monitoring_job( self, request: job_service.DeleteModelDeploymentMonitoringJobRequest, @@ -726,12 +967,35 @@ async def post_delete_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_deployment_monitoring_job` interceptor runs + before the `post_delete_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_delete_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_deployment_monitoring_job` will be passed to + `post_delete_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_delete_nas_job( self, request: job_service.DeleteNasJobRequest, @@ -751,12 +1015,35 @@ async def post_delete_nas_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_delete_nas_job` interceptor runs + before the `post_delete_nas_job_with_metadata` interceptor. """ return response + async def post_delete_nas_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_delete_nas_job_with_metadata` + interceptor in new development instead of the `post_delete_nas_job` interceptor. + When both interceptors are used, this `post_delete_nas_job_with_metadata` interceptor runs after the + `post_delete_nas_job` interceptor. The (possibly modified) response returned by + `post_delete_nas_job` will be passed to + `post_delete_nas_job_with_metadata`. + """ + return response, metadata + async def pre_get_batch_prediction_job( self, request: job_service.GetBatchPredictionJobRequest, @@ -777,12 +1064,37 @@ async def post_get_batch_prediction_job( ) -> batch_prediction_job.BatchPredictionJob: """Post-rpc interceptor for get_batch_prediction_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_batch_prediction_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_batch_prediction_job` interceptor runs + before the `post_get_batch_prediction_job_with_metadata` interceptor. """ return response + async def post_get_batch_prediction_job_with_metadata( + self, + response: batch_prediction_job.BatchPredictionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + batch_prediction_job.BatchPredictionJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_batch_prediction_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_batch_prediction_job_with_metadata` + interceptor in new development instead of the `post_get_batch_prediction_job` interceptor. + When both interceptors are used, this `post_get_batch_prediction_job_with_metadata` interceptor runs after the + `post_get_batch_prediction_job` interceptor. The (possibly modified) response returned by + `post_get_batch_prediction_job` will be passed to + `post_get_batch_prediction_job_with_metadata`. + """ + return response, metadata + async def pre_get_custom_job( self, request: job_service.GetCustomJobRequest, @@ -802,12 +1114,35 @@ async def post_get_custom_job( ) -> custom_job.CustomJob: """Post-rpc interceptor for get_custom_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_custom_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_custom_job` interceptor runs + before the `post_get_custom_job_with_metadata` interceptor. """ return response + async def post_get_custom_job_with_metadata( + self, + response: custom_job.CustomJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[custom_job.CustomJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_custom_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_custom_job_with_metadata` + interceptor in new development instead of the `post_get_custom_job` interceptor. + When both interceptors are used, this `post_get_custom_job_with_metadata` interceptor runs after the + `post_get_custom_job` interceptor. The (possibly modified) response returned by + `post_get_custom_job` will be passed to + `post_get_custom_job_with_metadata`. + """ + return response, metadata + async def pre_get_data_labeling_job( self, request: job_service.GetDataLabelingJobRequest, @@ -827,12 +1162,37 @@ async def post_get_data_labeling_job( ) -> data_labeling_job.DataLabelingJob: """Post-rpc interceptor for get_data_labeling_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_data_labeling_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_data_labeling_job` interceptor runs + before the `post_get_data_labeling_job_with_metadata` interceptor. """ return response + async def post_get_data_labeling_job_with_metadata( + self, + response: data_labeling_job.DataLabelingJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + data_labeling_job.DataLabelingJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_data_labeling_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_data_labeling_job_with_metadata` + interceptor in new development instead of the `post_get_data_labeling_job` interceptor. + When both interceptors are used, this `post_get_data_labeling_job_with_metadata` interceptor runs after the + `post_get_data_labeling_job` interceptor. The (possibly modified) response returned by + `post_get_data_labeling_job` will be passed to + `post_get_data_labeling_job_with_metadata`. + """ + return response, metadata + async def pre_get_hyperparameter_tuning_job( self, request: job_service.GetHyperparameterTuningJobRequest, @@ -853,12 +1213,38 @@ async def post_get_hyperparameter_tuning_job( ) -> hyperparameter_tuning_job.HyperparameterTuningJob: """Post-rpc interceptor for get_hyperparameter_tuning_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_hyperparameter_tuning_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_hyperparameter_tuning_job` interceptor runs + before the `post_get_hyperparameter_tuning_job_with_metadata` interceptor. """ return response + async def post_get_hyperparameter_tuning_job_with_metadata( + self, + response: hyperparameter_tuning_job.HyperparameterTuningJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + hyperparameter_tuning_job.HyperparameterTuningJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_hyperparameter_tuning_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_hyperparameter_tuning_job_with_metadata` + interceptor in new development instead of the `post_get_hyperparameter_tuning_job` interceptor. + When both interceptors are used, this `post_get_hyperparameter_tuning_job_with_metadata` interceptor runs after the + `post_get_hyperparameter_tuning_job` interceptor. The (possibly modified) response returned by + `post_get_hyperparameter_tuning_job` will be passed to + `post_get_hyperparameter_tuning_job_with_metadata`. + """ + return response, metadata + async def pre_get_model_deployment_monitoring_job( self, request: job_service.GetModelDeploymentMonitoringJobRequest, @@ -879,12 +1265,38 @@ async def post_get_model_deployment_monitoring_job( ) -> model_deployment_monitoring_job.ModelDeploymentMonitoringJob: """Post-rpc interceptor for get_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_deployment_monitoring_job` interceptor runs + before the `post_get_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_get_model_deployment_monitoring_job_with_metadata( + self, + response: model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_deployment_monitoring_job.ModelDeploymentMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_deployment_monitoring_job` will be passed to + `post_get_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_get_nas_job( self, request: job_service.GetNasJobRequest, @@ -900,12 +1312,35 @@ async def pre_get_nas_job( async def post_get_nas_job(self, response: nas_job.NasJob) -> nas_job.NasJob: """Post-rpc interceptor for get_nas_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_job` interceptor runs + before the `post_get_nas_job_with_metadata` interceptor. """ return response + async def post_get_nas_job_with_metadata( + self, + response: nas_job.NasJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_job_with_metadata` + interceptor in new development instead of the `post_get_nas_job` interceptor. + When both interceptors are used, this `post_get_nas_job_with_metadata` interceptor runs after the + `post_get_nas_job` interceptor. The (possibly modified) response returned by + `post_get_nas_job` will be passed to + `post_get_nas_job_with_metadata`. + """ + return response, metadata + async def pre_get_nas_trial_detail( self, request: job_service.GetNasTrialDetailRequest, @@ -925,12 +1360,35 @@ async def post_get_nas_trial_detail( ) -> nas_job.NasTrialDetail: """Post-rpc interceptor for get_nas_trial_detail - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_nas_trial_detail_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_get_nas_trial_detail` interceptor runs + before the `post_get_nas_trial_detail_with_metadata` interceptor. """ return response + async def post_get_nas_trial_detail_with_metadata( + self, + response: nas_job.NasTrialDetail, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[nas_job.NasTrialDetail, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_nas_trial_detail + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_get_nas_trial_detail_with_metadata` + interceptor in new development instead of the `post_get_nas_trial_detail` interceptor. + When both interceptors are used, this `post_get_nas_trial_detail_with_metadata` interceptor runs after the + `post_get_nas_trial_detail` interceptor. The (possibly modified) response returned by + `post_get_nas_trial_detail` will be passed to + `post_get_nas_trial_detail_with_metadata`. + """ + return response, metadata + async def pre_list_batch_prediction_jobs( self, request: job_service.ListBatchPredictionJobsRequest, @@ -951,12 +1409,38 @@ async def post_list_batch_prediction_jobs( ) -> job_service.ListBatchPredictionJobsResponse: """Post-rpc interceptor for list_batch_prediction_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_batch_prediction_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_batch_prediction_jobs` interceptor runs + before the `post_list_batch_prediction_jobs_with_metadata` interceptor. """ return response + async def post_list_batch_prediction_jobs_with_metadata( + self, + response: job_service.ListBatchPredictionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListBatchPredictionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_batch_prediction_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_batch_prediction_jobs_with_metadata` + interceptor in new development instead of the `post_list_batch_prediction_jobs` interceptor. + When both interceptors are used, this `post_list_batch_prediction_jobs_with_metadata` interceptor runs after the + `post_list_batch_prediction_jobs` interceptor. The (possibly modified) response returned by + `post_list_batch_prediction_jobs` will be passed to + `post_list_batch_prediction_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_custom_jobs( self, request: job_service.ListCustomJobsRequest, @@ -976,12 +1460,37 @@ async def post_list_custom_jobs( ) -> job_service.ListCustomJobsResponse: """Post-rpc interceptor for list_custom_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_custom_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_custom_jobs` interceptor runs + before the `post_list_custom_jobs_with_metadata` interceptor. """ return response + async def post_list_custom_jobs_with_metadata( + self, + response: job_service.ListCustomJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListCustomJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_custom_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_custom_jobs_with_metadata` + interceptor in new development instead of the `post_list_custom_jobs` interceptor. + When both interceptors are used, this `post_list_custom_jobs_with_metadata` interceptor runs after the + `post_list_custom_jobs` interceptor. The (possibly modified) response returned by + `post_list_custom_jobs` will be passed to + `post_list_custom_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_data_labeling_jobs( self, request: job_service.ListDataLabelingJobsRequest, @@ -1001,12 +1510,38 @@ async def post_list_data_labeling_jobs( ) -> job_service.ListDataLabelingJobsResponse: """Post-rpc interceptor for list_data_labeling_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_data_labeling_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_data_labeling_jobs` interceptor runs + before the `post_list_data_labeling_jobs_with_metadata` interceptor. """ return response + async def post_list_data_labeling_jobs_with_metadata( + self, + response: job_service.ListDataLabelingJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListDataLabelingJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_data_labeling_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_data_labeling_jobs_with_metadata` + interceptor in new development instead of the `post_list_data_labeling_jobs` interceptor. + When both interceptors are used, this `post_list_data_labeling_jobs_with_metadata` interceptor runs after the + `post_list_data_labeling_jobs` interceptor. The (possibly modified) response returned by + `post_list_data_labeling_jobs` will be passed to + `post_list_data_labeling_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_hyperparameter_tuning_jobs( self, request: job_service.ListHyperparameterTuningJobsRequest, @@ -1027,12 +1562,38 @@ async def post_list_hyperparameter_tuning_jobs( ) -> job_service.ListHyperparameterTuningJobsResponse: """Post-rpc interceptor for list_hyperparameter_tuning_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_hyperparameter_tuning_jobs` interceptor runs + before the `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor. """ return response + async def post_list_hyperparameter_tuning_jobs_with_metadata( + self, + response: job_service.ListHyperparameterTuningJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListHyperparameterTuningJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_hyperparameter_tuning_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_hyperparameter_tuning_jobs_with_metadata` + interceptor in new development instead of the `post_list_hyperparameter_tuning_jobs` interceptor. + When both interceptors are used, this `post_list_hyperparameter_tuning_jobs_with_metadata` interceptor runs after the + `post_list_hyperparameter_tuning_jobs` interceptor. The (possibly modified) response returned by + `post_list_hyperparameter_tuning_jobs` will be passed to + `post_list_hyperparameter_tuning_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_model_deployment_monitoring_jobs( self, request: job_service.ListModelDeploymentMonitoringJobsRequest, @@ -1053,12 +1614,38 @@ async def post_list_model_deployment_monitoring_jobs( ) -> job_service.ListModelDeploymentMonitoringJobsResponse: """Post-rpc interceptor for list_model_deployment_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_deployment_monitoring_jobs` interceptor runs + before the `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor. """ return response + async def post_list_model_deployment_monitoring_jobs_with_metadata( + self, + response: job_service.ListModelDeploymentMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListModelDeploymentMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_deployment_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_model_deployment_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_deployment_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_deployment_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_deployment_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_deployment_monitoring_jobs` will be passed to + `post_list_model_deployment_monitoring_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_nas_jobs( self, request: job_service.ListNasJobsRequest, @@ -1076,12 +1663,37 @@ async def post_list_nas_jobs( ) -> job_service.ListNasJobsResponse: """Post-rpc interceptor for list_nas_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_jobs` interceptor runs + before the `post_list_nas_jobs_with_metadata` interceptor. """ return response + async def post_list_nas_jobs_with_metadata( + self, + response: job_service.ListNasJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasJobsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_jobs_with_metadata` + interceptor in new development instead of the `post_list_nas_jobs` interceptor. + When both interceptors are used, this `post_list_nas_jobs_with_metadata` interceptor runs after the + `post_list_nas_jobs` interceptor. The (possibly modified) response returned by + `post_list_nas_jobs` will be passed to + `post_list_nas_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_nas_trial_details( self, request: job_service.ListNasTrialDetailsRequest, @@ -1101,12 +1713,37 @@ async def post_list_nas_trial_details( ) -> job_service.ListNasTrialDetailsResponse: """Post-rpc interceptor for list_nas_trial_details - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_nas_trial_details_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_list_nas_trial_details` interceptor runs + before the `post_list_nas_trial_details_with_metadata` interceptor. """ return response + async def post_list_nas_trial_details_with_metadata( + self, + response: job_service.ListNasTrialDetailsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.ListNasTrialDetailsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_nas_trial_details + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_list_nas_trial_details_with_metadata` + interceptor in new development instead of the `post_list_nas_trial_details` interceptor. + When both interceptors are used, this `post_list_nas_trial_details_with_metadata` interceptor runs after the + `post_list_nas_trial_details` interceptor. The (possibly modified) response returned by + `post_list_nas_trial_details` will be passed to + `post_list_nas_trial_details_with_metadata`. + """ + return response, metadata + async def pre_pause_model_deployment_monitoring_job( self, request: job_service.PauseModelDeploymentMonitoringJobRequest, @@ -1158,12 +1795,38 @@ async def post_search_model_deployment_monitoring_stats_anomalies( ) -> job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse: """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_deployment_monitoring_stats_anomalies` interceptor runs + before the `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor. """ return response + async def post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + self, + response: job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_deployment_monitoring_stats_anomalies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` + interceptor in new development instead of the `post_search_model_deployment_monitoring_stats_anomalies` interceptor. + When both interceptors are used, this `post_search_model_deployment_monitoring_stats_anomalies_with_metadata` interceptor runs after the + `post_search_model_deployment_monitoring_stats_anomalies` interceptor. The (possibly modified) response returned by + `post_search_model_deployment_monitoring_stats_anomalies` will be passed to + `post_search_model_deployment_monitoring_stats_anomalies_with_metadata`. + """ + return response, metadata + async def pre_update_model_deployment_monitoring_job( self, request: job_service.UpdateModelDeploymentMonitoringJobRequest, @@ -1184,12 +1847,35 @@ async def post_update_model_deployment_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_deployment_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_deployment_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the JobService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_deployment_monitoring_job` interceptor runs + before the `post_update_model_deployment_monitoring_job_with_metadata` interceptor. """ return response + async def post_update_model_deployment_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_deployment_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the JobService server but before it is returned to user code. + + We recommend only using this `post_update_model_deployment_monitoring_job_with_metadata` + interceptor in new development instead of the `post_update_model_deployment_monitoring_job` interceptor. + When both interceptors are used, this `post_update_model_deployment_monitoring_job_with_metadata` interceptor runs after the + `post_update_model_deployment_monitoring_job` interceptor. The (possibly modified) response returned by + `post_update_model_deployment_monitoring_job` will be passed to + `post_update_model_deployment_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2505,6 +3191,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2673,6 +3366,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2834,6 +3531,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3001,6 +3705,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3179,6 +3890,13 @@ async def __call__( resp = await self._interceptor.post_create_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3345,6 +4063,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3500,6 +4222,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3655,6 +4384,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3810,6 +4543,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3968,6 +4708,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4130,6 +4877,13 @@ async def __call__( resp = await self._interceptor.post_delete_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4287,6 +5041,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4446,6 +5204,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_batch_prediction_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_batch_prediction_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4610,6 +5375,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_custom_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_custom_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4765,6 +5534,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_data_labeling_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_data_labeling_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4926,6 +5699,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_hyperparameter_tuning_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_hyperparameter_tuning_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5094,6 +5874,13 @@ async def __call__( resp = await self._interceptor.post_get_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5252,6 +6039,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_nas_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_nas_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5408,6 +6199,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_nas_trial_detail(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_nas_trial_detail_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5562,6 +6357,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_batch_prediction_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_batch_prediction_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5718,6 +6520,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_custom_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_custom_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5874,6 +6680,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_data_labeling_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_data_labeling_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6033,6 +6846,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_hyperparameter_tuning_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_hyperparameter_tuning_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6198,6 +7018,13 @@ async def __call__( resp = await self._interceptor.post_list_model_deployment_monitoring_jobs( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_deployment_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6358,6 +7185,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_nas_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_nas_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6512,6 +7343,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_nas_trial_details(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_nas_trial_details_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6943,6 +7778,13 @@ async def __call__( resp = await self._interceptor.post_search_model_deployment_monitoring_stats_anomalies( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_model_deployment_monitoring_stats_anomalies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7113,6 +7955,13 @@ async def __call__( resp = await self._interceptor.post_update_model_deployment_monitoring_job( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_model_deployment_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/client.py b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/client.py index d01a5f351c..b6e1cf8af4 100644 --- a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -501,6 +503,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -904,16 +933,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -959,16 +992,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1131,16 +1168,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1252,16 +1293,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1374,16 +1419,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1434,16 +1483,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1489,16 +1542,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1544,16 +1601,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest.py index 51369f2ba5..495d03626b 100644 --- a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest.py @@ -112,12 +112,38 @@ def post_compute_tokens( ) -> llm_utility_service.ComputeTokensResponse: """Post-rpc interceptor for compute_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_compute_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_compute_tokens` interceptor runs + before the `post_compute_tokens_with_metadata` interceptor. """ return response + def post_compute_tokens_with_metadata( + self, + response: llm_utility_service.ComputeTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + llm_utility_service.ComputeTokensResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for compute_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_compute_tokens_with_metadata` + interceptor in new development instead of the `post_compute_tokens` interceptor. + When both interceptors are used, this `post_compute_tokens_with_metadata` interceptor runs after the + `post_compute_tokens` interceptor. The (possibly modified) response returned by + `post_compute_tokens` will be passed to + `post_compute_tokens_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -576,6 +602,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_compute_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_compute_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest_asyncio.py index 482aa238ae..13effcc566 100644 --- a/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/llm_utility_service/transports/rest_asyncio.py @@ -129,12 +129,38 @@ async def post_compute_tokens( ) -> llm_utility_service.ComputeTokensResponse: """Post-rpc interceptor for compute_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_compute_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the LlmUtilityService server but before - it is returned to user code. + it is returned to user code. This `post_compute_tokens` interceptor runs + before the `post_compute_tokens_with_metadata` interceptor. """ return response + async def post_compute_tokens_with_metadata( + self, + response: llm_utility_service.ComputeTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + llm_utility_service.ComputeTokensResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for compute_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the LlmUtilityService server but before it is returned to user code. + + We recommend only using this `post_compute_tokens_with_metadata` + interceptor in new development instead of the `post_compute_tokens` interceptor. + When both interceptors are used, this `post_compute_tokens_with_metadata` interceptor runs after the + `post_compute_tokens` interceptor. The (possibly modified) response returned by + `post_compute_tokens` will be passed to + `post_compute_tokens_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -646,6 +672,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_compute_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_compute_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/match_service/client.py b/google/cloud/aiplatform_v1beta1/services/match_service/client.py index eff1ed6d61..506d16e605 100644 --- a/google/cloud/aiplatform_v1beta1/services/match_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/match_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -501,6 +503,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -951,16 +980,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1006,16 +1039,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1178,16 +1215,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1299,16 +1340,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1421,16 +1466,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1481,16 +1530,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1536,16 +1589,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1591,16 +1648,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest.py index 9a6fc58b0e..20f349b620 100644 --- a/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest.py @@ -119,12 +119,37 @@ def post_find_neighbors( ) -> match_service.FindNeighborsResponse: """Post-rpc interceptor for find_neighbors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_find_neighbors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_find_neighbors` interceptor runs + before the `post_find_neighbors_with_metadata` interceptor. """ return response + def post_find_neighbors_with_metadata( + self, + response: match_service.FindNeighborsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.FindNeighborsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for find_neighbors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_find_neighbors_with_metadata` + interceptor in new development instead of the `post_find_neighbors` interceptor. + When both interceptors are used, this `post_find_neighbors_with_metadata` interceptor runs after the + `post_find_neighbors` interceptor. The (possibly modified) response returned by + `post_find_neighbors` will be passed to + `post_find_neighbors_with_metadata`. + """ + return response, metadata + def pre_read_index_datapoints( self, request: match_service.ReadIndexDatapointsRequest, @@ -145,12 +170,38 @@ def post_read_index_datapoints( ) -> match_service.ReadIndexDatapointsResponse: """Post-rpc interceptor for read_index_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_index_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_read_index_datapoints` interceptor runs + before the `post_read_index_datapoints_with_metadata` interceptor. """ return response + def post_read_index_datapoints_with_metadata( + self, + response: match_service.ReadIndexDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.ReadIndexDatapointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_index_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_read_index_datapoints_with_metadata` + interceptor in new development instead of the `post_read_index_datapoints` interceptor. + When both interceptors are used, this `post_read_index_datapoints_with_metadata` interceptor runs after the + `post_read_index_datapoints` interceptor. The (possibly modified) response returned by + `post_read_index_datapoints` will be passed to + `post_read_index_datapoints_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -609,6 +660,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_find_neighbors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_find_neighbors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -763,6 +818,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_index_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_index_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest_asyncio.py index e150d8d6a6..79506bc017 100644 --- a/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/match_service/transports/rest_asyncio.py @@ -136,12 +136,37 @@ async def post_find_neighbors( ) -> match_service.FindNeighborsResponse: """Post-rpc interceptor for find_neighbors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_find_neighbors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_find_neighbors` interceptor runs + before the `post_find_neighbors_with_metadata` interceptor. """ return response + async def post_find_neighbors_with_metadata( + self, + response: match_service.FindNeighborsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.FindNeighborsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for find_neighbors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_find_neighbors_with_metadata` + interceptor in new development instead of the `post_find_neighbors` interceptor. + When both interceptors are used, this `post_find_neighbors_with_metadata` interceptor runs after the + `post_find_neighbors` interceptor. The (possibly modified) response returned by + `post_find_neighbors` will be passed to + `post_find_neighbors_with_metadata`. + """ + return response, metadata + async def pre_read_index_datapoints( self, request: match_service.ReadIndexDatapointsRequest, @@ -162,12 +187,38 @@ async def post_read_index_datapoints( ) -> match_service.ReadIndexDatapointsResponse: """Post-rpc interceptor for read_index_datapoints - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_index_datapoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MatchService server but before - it is returned to user code. + it is returned to user code. This `post_read_index_datapoints` interceptor runs + before the `post_read_index_datapoints_with_metadata` interceptor. """ return response + async def post_read_index_datapoints_with_metadata( + self, + response: match_service.ReadIndexDatapointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + match_service.ReadIndexDatapointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_index_datapoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MatchService server but before it is returned to user code. + + We recommend only using this `post_read_index_datapoints_with_metadata` + interceptor in new development instead of the `post_read_index_datapoints` interceptor. + When both interceptors are used, this `post_read_index_datapoints_with_metadata` interceptor runs after the + `post_read_index_datapoints` interceptor. The (possibly modified) response returned by + `post_read_index_datapoints` will be passed to + `post_read_index_datapoints_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -684,6 +735,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_find_neighbors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_find_neighbors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -847,6 +902,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_index_datapoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_index_datapoints_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/metadata_service/client.py b/google/cloud/aiplatform_v1beta1/services/metadata_service/client.py index 5f4263d71f..1c1a927861 100644 --- a/google/cloud/aiplatform_v1beta1/services/metadata_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/metadata_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -617,6 +619,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -4805,16 +4834,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -4860,16 +4893,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -5032,16 +5069,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5153,16 +5194,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5275,16 +5320,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5335,16 +5384,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -5390,16 +5443,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -5445,16 +5502,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest.py index 22f073ffba..6536731c33 100644 --- a/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest.py @@ -371,12 +371,38 @@ def post_add_context_artifacts_and_executions( ) -> metadata_service.AddContextArtifactsAndExecutionsResponse: """Post-rpc interceptor for add_context_artifacts_and_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_artifacts_and_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_artifacts_and_executions` interceptor runs + before the `post_add_context_artifacts_and_executions_with_metadata` interceptor. """ return response + def post_add_context_artifacts_and_executions_with_metadata( + self, + response: metadata_service.AddContextArtifactsAndExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextArtifactsAndExecutionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_artifacts_and_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_artifacts_and_executions_with_metadata` + interceptor in new development instead of the `post_add_context_artifacts_and_executions` interceptor. + When both interceptors are used, this `post_add_context_artifacts_and_executions_with_metadata` interceptor runs after the + `post_add_context_artifacts_and_executions` interceptor. The (possibly modified) response returned by + `post_add_context_artifacts_and_executions` will be passed to + `post_add_context_artifacts_and_executions_with_metadata`. + """ + return response, metadata + def pre_add_context_children( self, request: metadata_service.AddContextChildrenRequest, @@ -397,12 +423,38 @@ def post_add_context_children( ) -> metadata_service.AddContextChildrenResponse: """Post-rpc interceptor for add_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_children` interceptor runs + before the `post_add_context_children_with_metadata` interceptor. """ return response + def post_add_context_children_with_metadata( + self, + response: metadata_service.AddContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_children_with_metadata` + interceptor in new development instead of the `post_add_context_children` interceptor. + When both interceptors are used, this `post_add_context_children_with_metadata` interceptor runs after the + `post_add_context_children` interceptor. The (possibly modified) response returned by + `post_add_context_children` will be passed to + `post_add_context_children_with_metadata`. + """ + return response, metadata + def pre_add_execution_events( self, request: metadata_service.AddExecutionEventsRequest, @@ -423,12 +475,38 @@ def post_add_execution_events( ) -> metadata_service.AddExecutionEventsResponse: """Post-rpc interceptor for add_execution_events - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_execution_events_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_execution_events` interceptor runs + before the `post_add_execution_events_with_metadata` interceptor. """ return response + def post_add_execution_events_with_metadata( + self, + response: metadata_service.AddExecutionEventsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddExecutionEventsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_execution_events + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_execution_events_with_metadata` + interceptor in new development instead of the `post_add_execution_events` interceptor. + When both interceptors are used, this `post_add_execution_events_with_metadata` interceptor runs after the + `post_add_execution_events` interceptor. The (possibly modified) response returned by + `post_add_execution_events` will be passed to + `post_add_execution_events_with_metadata`. + """ + return response, metadata + def pre_create_artifact( self, request: metadata_service.CreateArtifactRequest, @@ -448,12 +526,35 @@ def post_create_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for create_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_artifact` interceptor runs + before the `post_create_artifact_with_metadata` interceptor. """ return response + def post_create_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_artifact_with_metadata` + interceptor in new development instead of the `post_create_artifact` interceptor. + When both interceptors are used, this `post_create_artifact_with_metadata` interceptor runs after the + `post_create_artifact` interceptor. The (possibly modified) response returned by + `post_create_artifact` will be passed to + `post_create_artifact_with_metadata`. + """ + return response, metadata + def pre_create_context( self, request: metadata_service.CreateContextRequest, @@ -471,12 +572,35 @@ def pre_create_context( def post_create_context(self, response: gca_context.Context) -> gca_context.Context: """Post-rpc interceptor for create_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_context` interceptor runs + before the `post_create_context_with_metadata` interceptor. """ return response + def post_create_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_context_with_metadata` + interceptor in new development instead of the `post_create_context` interceptor. + When both interceptors are used, this `post_create_context_with_metadata` interceptor runs after the + `post_create_context` interceptor. The (possibly modified) response returned by + `post_create_context` will be passed to + `post_create_context_with_metadata`. + """ + return response, metadata + def pre_create_execution( self, request: metadata_service.CreateExecutionRequest, @@ -496,12 +620,35 @@ def post_create_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for create_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_execution` interceptor runs + before the `post_create_execution_with_metadata` interceptor. """ return response + def post_create_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_execution_with_metadata` + interceptor in new development instead of the `post_create_execution` interceptor. + When both interceptors are used, this `post_create_execution_with_metadata` interceptor runs after the + `post_create_execution` interceptor. The (possibly modified) response returned by + `post_create_execution` will be passed to + `post_create_execution_with_metadata`. + """ + return response, metadata + def pre_create_metadata_schema( self, request: metadata_service.CreateMetadataSchemaRequest, @@ -522,12 +669,37 @@ def post_create_metadata_schema( ) -> gca_metadata_schema.MetadataSchema: """Post-rpc interceptor for create_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_schema` interceptor runs + before the `post_create_metadata_schema_with_metadata` interceptor. """ return response + def post_create_metadata_schema_with_metadata( + self, + response: gca_metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_schema_with_metadata` + interceptor in new development instead of the `post_create_metadata_schema` interceptor. + When both interceptors are used, this `post_create_metadata_schema_with_metadata` interceptor runs after the + `post_create_metadata_schema` interceptor. The (possibly modified) response returned by + `post_create_metadata_schema` will be passed to + `post_create_metadata_schema_with_metadata`. + """ + return response, metadata + def pre_create_metadata_store( self, request: metadata_service.CreateMetadataStoreRequest, @@ -548,12 +720,35 @@ def post_create_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_store` interceptor runs + before the `post_create_metadata_store_with_metadata` interceptor. """ return response + def post_create_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_store_with_metadata` + interceptor in new development instead of the `post_create_metadata_store` interceptor. + When both interceptors are used, this `post_create_metadata_store_with_metadata` interceptor runs after the + `post_create_metadata_store` interceptor. The (possibly modified) response returned by + `post_create_metadata_store` will be passed to + `post_create_metadata_store_with_metadata`. + """ + return response, metadata + def pre_delete_artifact( self, request: metadata_service.DeleteArtifactRequest, @@ -573,12 +768,35 @@ def post_delete_artifact( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_artifact` interceptor runs + before the `post_delete_artifact_with_metadata` interceptor. """ return response + def post_delete_artifact_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_artifact_with_metadata` + interceptor in new development instead of the `post_delete_artifact` interceptor. + When both interceptors are used, this `post_delete_artifact_with_metadata` interceptor runs after the + `post_delete_artifact` interceptor. The (possibly modified) response returned by + `post_delete_artifact` will be passed to + `post_delete_artifact_with_metadata`. + """ + return response, metadata + def pre_delete_context( self, request: metadata_service.DeleteContextRequest, @@ -598,12 +816,35 @@ def post_delete_context( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_context` interceptor runs + before the `post_delete_context_with_metadata` interceptor. """ return response + def post_delete_context_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_context_with_metadata` + interceptor in new development instead of the `post_delete_context` interceptor. + When both interceptors are used, this `post_delete_context_with_metadata` interceptor runs after the + `post_delete_context` interceptor. The (possibly modified) response returned by + `post_delete_context` will be passed to + `post_delete_context_with_metadata`. + """ + return response, metadata + def pre_delete_execution( self, request: metadata_service.DeleteExecutionRequest, @@ -623,12 +864,35 @@ def post_delete_execution( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_execution` interceptor runs + before the `post_delete_execution_with_metadata` interceptor. """ return response + def post_delete_execution_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_execution_with_metadata` + interceptor in new development instead of the `post_delete_execution` interceptor. + When both interceptors are used, this `post_delete_execution_with_metadata` interceptor runs after the + `post_delete_execution` interceptor. The (possibly modified) response returned by + `post_delete_execution` will be passed to + `post_delete_execution_with_metadata`. + """ + return response, metadata + def pre_delete_metadata_store( self, request: metadata_service.DeleteMetadataStoreRequest, @@ -649,12 +913,35 @@ def post_delete_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_metadata_store` interceptor runs + before the `post_delete_metadata_store_with_metadata` interceptor. """ return response + def post_delete_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_metadata_store_with_metadata` + interceptor in new development instead of the `post_delete_metadata_store` interceptor. + When both interceptors are used, this `post_delete_metadata_store_with_metadata` interceptor runs after the + `post_delete_metadata_store` interceptor. The (possibly modified) response returned by + `post_delete_metadata_store` will be passed to + `post_delete_metadata_store_with_metadata`. + """ + return response, metadata + def pre_get_artifact( self, request: metadata_service.GetArtifactRequest, @@ -672,12 +959,35 @@ def pre_get_artifact( def post_get_artifact(self, response: artifact.Artifact) -> artifact.Artifact: """Post-rpc interceptor for get_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_artifact` interceptor runs + before the `post_get_artifact_with_metadata` interceptor. """ return response + def post_get_artifact_with_metadata( + self, + response: artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_artifact_with_metadata` + interceptor in new development instead of the `post_get_artifact` interceptor. + When both interceptors are used, this `post_get_artifact_with_metadata` interceptor runs after the + `post_get_artifact` interceptor. The (possibly modified) response returned by + `post_get_artifact` will be passed to + `post_get_artifact_with_metadata`. + """ + return response, metadata + def pre_get_context( self, request: metadata_service.GetContextRequest, @@ -695,12 +1005,35 @@ def pre_get_context( def post_get_context(self, response: context.Context) -> context.Context: """Post-rpc interceptor for get_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_context` interceptor runs + before the `post_get_context_with_metadata` interceptor. """ return response + def post_get_context_with_metadata( + self, + response: context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_context_with_metadata` + interceptor in new development instead of the `post_get_context` interceptor. + When both interceptors are used, this `post_get_context_with_metadata` interceptor runs after the + `post_get_context` interceptor. The (possibly modified) response returned by + `post_get_context` will be passed to + `post_get_context_with_metadata`. + """ + return response, metadata + def pre_get_execution( self, request: metadata_service.GetExecutionRequest, @@ -718,12 +1051,35 @@ def pre_get_execution( def post_get_execution(self, response: execution.Execution) -> execution.Execution: """Post-rpc interceptor for get_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_execution` interceptor runs + before the `post_get_execution_with_metadata` interceptor. """ return response + def post_get_execution_with_metadata( + self, + response: execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_execution_with_metadata` + interceptor in new development instead of the `post_get_execution` interceptor. + When both interceptors are used, this `post_get_execution_with_metadata` interceptor runs after the + `post_get_execution` interceptor. The (possibly modified) response returned by + `post_get_execution` will be passed to + `post_get_execution_with_metadata`. + """ + return response, metadata + def pre_get_metadata_schema( self, request: metadata_service.GetMetadataSchemaRequest, @@ -744,12 +1100,35 @@ def post_get_metadata_schema( ) -> metadata_schema.MetadataSchema: """Post-rpc interceptor for get_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_schema` interceptor runs + before the `post_get_metadata_schema_with_metadata` interceptor. """ return response + def post_get_metadata_schema_with_metadata( + self, + response: metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_schema_with_metadata` + interceptor in new development instead of the `post_get_metadata_schema` interceptor. + When both interceptors are used, this `post_get_metadata_schema_with_metadata` interceptor runs after the + `post_get_metadata_schema` interceptor. The (possibly modified) response returned by + `post_get_metadata_schema` will be passed to + `post_get_metadata_schema_with_metadata`. + """ + return response, metadata + def pre_get_metadata_store( self, request: metadata_service.GetMetadataStoreRequest, @@ -770,12 +1149,35 @@ def post_get_metadata_store( ) -> metadata_store.MetadataStore: """Post-rpc interceptor for get_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_store` interceptor runs + before the `post_get_metadata_store_with_metadata` interceptor. """ return response + def post_get_metadata_store_with_metadata( + self, + response: metadata_store.MetadataStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_store.MetadataStore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_store_with_metadata` + interceptor in new development instead of the `post_get_metadata_store` interceptor. + When both interceptors are used, this `post_get_metadata_store_with_metadata` interceptor runs after the + `post_get_metadata_store` interceptor. The (possibly modified) response returned by + `post_get_metadata_store` will be passed to + `post_get_metadata_store_with_metadata`. + """ + return response, metadata + def pre_list_artifacts( self, request: metadata_service.ListArtifactsRequest, @@ -795,12 +1197,37 @@ def post_list_artifacts( ) -> metadata_service.ListArtifactsResponse: """Post-rpc interceptor for list_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_artifacts` interceptor runs + before the `post_list_artifacts_with_metadata` interceptor. """ return response + def post_list_artifacts_with_metadata( + self, + response: metadata_service.ListArtifactsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListArtifactsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_artifacts_with_metadata` + interceptor in new development instead of the `post_list_artifacts` interceptor. + When both interceptors are used, this `post_list_artifacts_with_metadata` interceptor runs after the + `post_list_artifacts` interceptor. The (possibly modified) response returned by + `post_list_artifacts` will be passed to + `post_list_artifacts_with_metadata`. + """ + return response, metadata + def pre_list_contexts( self, request: metadata_service.ListContextsRequest, @@ -820,12 +1247,37 @@ def post_list_contexts( ) -> metadata_service.ListContextsResponse: """Post-rpc interceptor for list_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_contexts` interceptor runs + before the `post_list_contexts_with_metadata` interceptor. """ return response + def post_list_contexts_with_metadata( + self, + response: metadata_service.ListContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListContextsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_contexts_with_metadata` + interceptor in new development instead of the `post_list_contexts` interceptor. + When both interceptors are used, this `post_list_contexts_with_metadata` interceptor runs after the + `post_list_contexts` interceptor. The (possibly modified) response returned by + `post_list_contexts` will be passed to + `post_list_contexts_with_metadata`. + """ + return response, metadata + def pre_list_executions( self, request: metadata_service.ListExecutionsRequest, @@ -845,12 +1297,37 @@ def post_list_executions( ) -> metadata_service.ListExecutionsResponse: """Post-rpc interceptor for list_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_executions` interceptor runs + before the `post_list_executions_with_metadata` interceptor. """ return response + def post_list_executions_with_metadata( + self, + response: metadata_service.ListExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListExecutionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_executions_with_metadata` + interceptor in new development instead of the `post_list_executions` interceptor. + When both interceptors are used, this `post_list_executions_with_metadata` interceptor runs after the + `post_list_executions` interceptor. The (possibly modified) response returned by + `post_list_executions` will be passed to + `post_list_executions_with_metadata`. + """ + return response, metadata + def pre_list_metadata_schemas( self, request: metadata_service.ListMetadataSchemasRequest, @@ -871,12 +1348,38 @@ def post_list_metadata_schemas( ) -> metadata_service.ListMetadataSchemasResponse: """Post-rpc interceptor for list_metadata_schemas - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_schemas_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_schemas` interceptor runs + before the `post_list_metadata_schemas_with_metadata` interceptor. """ return response + def post_list_metadata_schemas_with_metadata( + self, + response: metadata_service.ListMetadataSchemasResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataSchemasResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_schemas + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_schemas_with_metadata` + interceptor in new development instead of the `post_list_metadata_schemas` interceptor. + When both interceptors are used, this `post_list_metadata_schemas_with_metadata` interceptor runs after the + `post_list_metadata_schemas` interceptor. The (possibly modified) response returned by + `post_list_metadata_schemas` will be passed to + `post_list_metadata_schemas_with_metadata`. + """ + return response, metadata + def pre_list_metadata_stores( self, request: metadata_service.ListMetadataStoresRequest, @@ -897,12 +1400,38 @@ def post_list_metadata_stores( ) -> metadata_service.ListMetadataStoresResponse: """Post-rpc interceptor for list_metadata_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_stores` interceptor runs + before the `post_list_metadata_stores_with_metadata` interceptor. """ return response + def post_list_metadata_stores_with_metadata( + self, + response: metadata_service.ListMetadataStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_stores_with_metadata` + interceptor in new development instead of the `post_list_metadata_stores` interceptor. + When both interceptors are used, this `post_list_metadata_stores_with_metadata` interceptor runs after the + `post_list_metadata_stores` interceptor. The (possibly modified) response returned by + `post_list_metadata_stores` will be passed to + `post_list_metadata_stores_with_metadata`. + """ + return response, metadata + def pre_purge_artifacts( self, request: metadata_service.PurgeArtifactsRequest, @@ -922,12 +1451,35 @@ def post_purge_artifacts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_artifacts` interceptor runs + before the `post_purge_artifacts_with_metadata` interceptor. """ return response + def post_purge_artifacts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_artifacts_with_metadata` + interceptor in new development instead of the `post_purge_artifacts` interceptor. + When both interceptors are used, this `post_purge_artifacts_with_metadata` interceptor runs after the + `post_purge_artifacts` interceptor. The (possibly modified) response returned by + `post_purge_artifacts` will be passed to + `post_purge_artifacts_with_metadata`. + """ + return response, metadata + def pre_purge_contexts( self, request: metadata_service.PurgeContextsRequest, @@ -947,12 +1499,35 @@ def post_purge_contexts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_contexts` interceptor runs + before the `post_purge_contexts_with_metadata` interceptor. """ return response + def post_purge_contexts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_contexts_with_metadata` + interceptor in new development instead of the `post_purge_contexts` interceptor. + When both interceptors are used, this `post_purge_contexts_with_metadata` interceptor runs after the + `post_purge_contexts` interceptor. The (possibly modified) response returned by + `post_purge_contexts` will be passed to + `post_purge_contexts_with_metadata`. + """ + return response, metadata + def pre_purge_executions( self, request: metadata_service.PurgeExecutionsRequest, @@ -972,12 +1547,35 @@ def post_purge_executions( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_executions` interceptor runs + before the `post_purge_executions_with_metadata` interceptor. """ return response + def post_purge_executions_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_executions_with_metadata` + interceptor in new development instead of the `post_purge_executions` interceptor. + When both interceptors are used, this `post_purge_executions_with_metadata` interceptor runs after the + `post_purge_executions` interceptor. The (possibly modified) response returned by + `post_purge_executions` will be passed to + `post_purge_executions_with_metadata`. + """ + return response, metadata + def pre_query_artifact_lineage_subgraph( self, request: metadata_service.QueryArtifactLineageSubgraphRequest, @@ -998,12 +1596,37 @@ def post_query_artifact_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_artifact_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_artifact_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_artifact_lineage_subgraph` interceptor runs + before the `post_query_artifact_lineage_subgraph_with_metadata` interceptor. """ return response + def post_query_artifact_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_artifact_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_artifact_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_artifact_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_artifact_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_artifact_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_artifact_lineage_subgraph` will be passed to + `post_query_artifact_lineage_subgraph_with_metadata`. + """ + return response, metadata + def pre_query_context_lineage_subgraph( self, request: metadata_service.QueryContextLineageSubgraphRequest, @@ -1024,12 +1647,37 @@ def post_query_context_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_context_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_context_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_context_lineage_subgraph` interceptor runs + before the `post_query_context_lineage_subgraph_with_metadata` interceptor. """ return response + def post_query_context_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_context_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_context_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_context_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_context_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_context_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_context_lineage_subgraph` will be passed to + `post_query_context_lineage_subgraph_with_metadata`. + """ + return response, metadata + def pre_query_execution_inputs_and_outputs( self, request: metadata_service.QueryExecutionInputsAndOutputsRequest, @@ -1050,12 +1698,37 @@ def post_query_execution_inputs_and_outputs( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_execution_inputs_and_outputs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_execution_inputs_and_outputs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_execution_inputs_and_outputs` interceptor runs + before the `post_query_execution_inputs_and_outputs_with_metadata` interceptor. """ return response + def post_query_execution_inputs_and_outputs_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_execution_inputs_and_outputs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_execution_inputs_and_outputs_with_metadata` + interceptor in new development instead of the `post_query_execution_inputs_and_outputs` interceptor. + When both interceptors are used, this `post_query_execution_inputs_and_outputs_with_metadata` interceptor runs after the + `post_query_execution_inputs_and_outputs` interceptor. The (possibly modified) response returned by + `post_query_execution_inputs_and_outputs` will be passed to + `post_query_execution_inputs_and_outputs_with_metadata`. + """ + return response, metadata + def pre_remove_context_children( self, request: metadata_service.RemoveContextChildrenRequest, @@ -1076,12 +1749,38 @@ def post_remove_context_children( ) -> metadata_service.RemoveContextChildrenResponse: """Post-rpc interceptor for remove_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_remove_context_children` interceptor runs + before the `post_remove_context_children_with_metadata` interceptor. """ return response + def post_remove_context_children_with_metadata( + self, + response: metadata_service.RemoveContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.RemoveContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for remove_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_remove_context_children_with_metadata` + interceptor in new development instead of the `post_remove_context_children` interceptor. + When both interceptors are used, this `post_remove_context_children_with_metadata` interceptor runs after the + `post_remove_context_children` interceptor. The (possibly modified) response returned by + `post_remove_context_children` will be passed to + `post_remove_context_children_with_metadata`. + """ + return response, metadata + def pre_update_artifact( self, request: metadata_service.UpdateArtifactRequest, @@ -1101,12 +1800,35 @@ def post_update_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for update_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_artifact` interceptor runs + before the `post_update_artifact_with_metadata` interceptor. """ return response + def post_update_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_artifact_with_metadata` + interceptor in new development instead of the `post_update_artifact` interceptor. + When both interceptors are used, this `post_update_artifact_with_metadata` interceptor runs after the + `post_update_artifact` interceptor. The (possibly modified) response returned by + `post_update_artifact` will be passed to + `post_update_artifact_with_metadata`. + """ + return response, metadata + def pre_update_context( self, request: metadata_service.UpdateContextRequest, @@ -1124,12 +1846,35 @@ def pre_update_context( def post_update_context(self, response: gca_context.Context) -> gca_context.Context: """Post-rpc interceptor for update_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_context` interceptor runs + before the `post_update_context_with_metadata` interceptor. """ return response + def post_update_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_context_with_metadata` + interceptor in new development instead of the `post_update_context` interceptor. + When both interceptors are used, this `post_update_context_with_metadata` interceptor runs after the + `post_update_context` interceptor. The (possibly modified) response returned by + `post_update_context` will be passed to + `post_update_context_with_metadata`. + """ + return response, metadata + def pre_update_execution( self, request: metadata_service.UpdateExecutionRequest, @@ -1149,12 +1894,35 @@ def post_update_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for update_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_execution` interceptor runs + before the `post_update_execution_with_metadata` interceptor. """ return response + def post_update_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_execution_with_metadata` + interceptor in new development instead of the `post_update_execution` interceptor. + When both interceptors are used, this `post_update_execution_with_metadata` interceptor runs after the + `post_update_execution` interceptor. The (possibly modified) response returned by + `post_update_execution` will be passed to + `post_update_execution_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3706,6 +4474,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_context_artifacts_and_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_add_context_artifacts_and_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3861,6 +4636,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4016,6 +4795,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_execution_events(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_execution_events_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4166,6 +4949,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4314,6 +5101,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4464,6 +5255,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4615,6 +5410,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4769,6 +5568,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4912,6 +5715,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5055,6 +5862,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5200,6 +6011,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5346,6 +6161,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5488,6 +6307,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5630,6 +6453,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5772,6 +6599,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5917,6 +6748,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6064,6 +6899,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6208,6 +7047,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6354,6 +7197,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6500,6 +7347,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6649,6 +7500,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_metadata_schemas(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_metadata_schemas_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6798,6 +7653,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_metadata_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_metadata_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6949,6 +7808,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7098,6 +7961,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7249,6 +8116,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_purge_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_purge_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7398,6 +8269,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_artifact_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_artifact_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7551,6 +8429,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_context_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_context_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7705,6 +8590,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_execution_inputs_and_outputs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_query_execution_inputs_and_outputs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7862,6 +8754,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_remove_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_remove_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8012,6 +8908,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8160,6 +9060,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8310,6 +9214,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest_asyncio.py index 54f3f59975..afaede4df6 100644 --- a/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/metadata_service/transports/rest_asyncio.py @@ -389,12 +389,38 @@ async def post_add_context_artifacts_and_executions( ) -> metadata_service.AddContextArtifactsAndExecutionsResponse: """Post-rpc interceptor for add_context_artifacts_and_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_artifacts_and_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_artifacts_and_executions` interceptor runs + before the `post_add_context_artifacts_and_executions_with_metadata` interceptor. """ return response + async def post_add_context_artifacts_and_executions_with_metadata( + self, + response: metadata_service.AddContextArtifactsAndExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextArtifactsAndExecutionsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_artifacts_and_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_artifacts_and_executions_with_metadata` + interceptor in new development instead of the `post_add_context_artifacts_and_executions` interceptor. + When both interceptors are used, this `post_add_context_artifacts_and_executions_with_metadata` interceptor runs after the + `post_add_context_artifacts_and_executions` interceptor. The (possibly modified) response returned by + `post_add_context_artifacts_and_executions` will be passed to + `post_add_context_artifacts_and_executions_with_metadata`. + """ + return response, metadata + async def pre_add_context_children( self, request: metadata_service.AddContextChildrenRequest, @@ -415,12 +441,38 @@ async def post_add_context_children( ) -> metadata_service.AddContextChildrenResponse: """Post-rpc interceptor for add_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_context_children` interceptor runs + before the `post_add_context_children_with_metadata` interceptor. """ return response + async def post_add_context_children_with_metadata( + self, + response: metadata_service.AddContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_context_children_with_metadata` + interceptor in new development instead of the `post_add_context_children` interceptor. + When both interceptors are used, this `post_add_context_children_with_metadata` interceptor runs after the + `post_add_context_children` interceptor. The (possibly modified) response returned by + `post_add_context_children` will be passed to + `post_add_context_children_with_metadata`. + """ + return response, metadata + async def pre_add_execution_events( self, request: metadata_service.AddExecutionEventsRequest, @@ -441,12 +493,38 @@ async def post_add_execution_events( ) -> metadata_service.AddExecutionEventsResponse: """Post-rpc interceptor for add_execution_events - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_execution_events_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_add_execution_events` interceptor runs + before the `post_add_execution_events_with_metadata` interceptor. """ return response + async def post_add_execution_events_with_metadata( + self, + response: metadata_service.AddExecutionEventsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.AddExecutionEventsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for add_execution_events + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_add_execution_events_with_metadata` + interceptor in new development instead of the `post_add_execution_events` interceptor. + When both interceptors are used, this `post_add_execution_events_with_metadata` interceptor runs after the + `post_add_execution_events` interceptor. The (possibly modified) response returned by + `post_add_execution_events` will be passed to + `post_add_execution_events_with_metadata`. + """ + return response, metadata + async def pre_create_artifact( self, request: metadata_service.CreateArtifactRequest, @@ -466,12 +544,35 @@ async def post_create_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for create_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_artifact` interceptor runs + before the `post_create_artifact_with_metadata` interceptor. """ return response + async def post_create_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_artifact_with_metadata` + interceptor in new development instead of the `post_create_artifact` interceptor. + When both interceptors are used, this `post_create_artifact_with_metadata` interceptor runs after the + `post_create_artifact` interceptor. The (possibly modified) response returned by + `post_create_artifact` will be passed to + `post_create_artifact_with_metadata`. + """ + return response, metadata + async def pre_create_context( self, request: metadata_service.CreateContextRequest, @@ -491,12 +592,35 @@ async def post_create_context( ) -> gca_context.Context: """Post-rpc interceptor for create_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_context` interceptor runs + before the `post_create_context_with_metadata` interceptor. """ return response + async def post_create_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_context_with_metadata` + interceptor in new development instead of the `post_create_context` interceptor. + When both interceptors are used, this `post_create_context_with_metadata` interceptor runs after the + `post_create_context` interceptor. The (possibly modified) response returned by + `post_create_context` will be passed to + `post_create_context_with_metadata`. + """ + return response, metadata + async def pre_create_execution( self, request: metadata_service.CreateExecutionRequest, @@ -516,12 +640,35 @@ async def post_create_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for create_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_execution` interceptor runs + before the `post_create_execution_with_metadata` interceptor. """ return response + async def post_create_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_execution_with_metadata` + interceptor in new development instead of the `post_create_execution` interceptor. + When both interceptors are used, this `post_create_execution_with_metadata` interceptor runs after the + `post_create_execution` interceptor. The (possibly modified) response returned by + `post_create_execution` will be passed to + `post_create_execution_with_metadata`. + """ + return response, metadata + async def pre_create_metadata_schema( self, request: metadata_service.CreateMetadataSchemaRequest, @@ -542,12 +689,37 @@ async def post_create_metadata_schema( ) -> gca_metadata_schema.MetadataSchema: """Post-rpc interceptor for create_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_schema` interceptor runs + before the `post_create_metadata_schema_with_metadata` interceptor. """ return response + async def post_create_metadata_schema_with_metadata( + self, + response: gca_metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_schema_with_metadata` + interceptor in new development instead of the `post_create_metadata_schema` interceptor. + When both interceptors are used, this `post_create_metadata_schema_with_metadata` interceptor runs after the + `post_create_metadata_schema` interceptor. The (possibly modified) response returned by + `post_create_metadata_schema` will be passed to + `post_create_metadata_schema_with_metadata`. + """ + return response, metadata + async def pre_create_metadata_store( self, request: metadata_service.CreateMetadataStoreRequest, @@ -568,12 +740,35 @@ async def post_create_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_create_metadata_store` interceptor runs + before the `post_create_metadata_store_with_metadata` interceptor. """ return response + async def post_create_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_create_metadata_store_with_metadata` + interceptor in new development instead of the `post_create_metadata_store` interceptor. + When both interceptors are used, this `post_create_metadata_store_with_metadata` interceptor runs after the + `post_create_metadata_store` interceptor. The (possibly modified) response returned by + `post_create_metadata_store` will be passed to + `post_create_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_delete_artifact( self, request: metadata_service.DeleteArtifactRequest, @@ -593,12 +788,35 @@ async def post_delete_artifact( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_artifact` interceptor runs + before the `post_delete_artifact_with_metadata` interceptor. """ return response + async def post_delete_artifact_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_artifact_with_metadata` + interceptor in new development instead of the `post_delete_artifact` interceptor. + When both interceptors are used, this `post_delete_artifact_with_metadata` interceptor runs after the + `post_delete_artifact` interceptor. The (possibly modified) response returned by + `post_delete_artifact` will be passed to + `post_delete_artifact_with_metadata`. + """ + return response, metadata + async def pre_delete_context( self, request: metadata_service.DeleteContextRequest, @@ -618,12 +836,35 @@ async def post_delete_context( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_context` interceptor runs + before the `post_delete_context_with_metadata` interceptor. """ return response + async def post_delete_context_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_context_with_metadata` + interceptor in new development instead of the `post_delete_context` interceptor. + When both interceptors are used, this `post_delete_context_with_metadata` interceptor runs after the + `post_delete_context` interceptor. The (possibly modified) response returned by + `post_delete_context` will be passed to + `post_delete_context_with_metadata`. + """ + return response, metadata + async def pre_delete_execution( self, request: metadata_service.DeleteExecutionRequest, @@ -643,12 +884,35 @@ async def post_delete_execution( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_execution` interceptor runs + before the `post_delete_execution_with_metadata` interceptor. """ return response + async def post_delete_execution_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_execution_with_metadata` + interceptor in new development instead of the `post_delete_execution` interceptor. + When both interceptors are used, this `post_delete_execution_with_metadata` interceptor runs after the + `post_delete_execution` interceptor. The (possibly modified) response returned by + `post_delete_execution` will be passed to + `post_delete_execution_with_metadata`. + """ + return response, metadata + async def pre_delete_metadata_store( self, request: metadata_service.DeleteMetadataStoreRequest, @@ -669,12 +933,35 @@ async def post_delete_metadata_store( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_metadata_store` interceptor runs + before the `post_delete_metadata_store_with_metadata` interceptor. """ return response + async def post_delete_metadata_store_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_delete_metadata_store_with_metadata` + interceptor in new development instead of the `post_delete_metadata_store` interceptor. + When both interceptors are used, this `post_delete_metadata_store_with_metadata` interceptor runs after the + `post_delete_metadata_store` interceptor. The (possibly modified) response returned by + `post_delete_metadata_store` will be passed to + `post_delete_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_get_artifact( self, request: metadata_service.GetArtifactRequest, @@ -692,12 +979,35 @@ async def pre_get_artifact( async def post_get_artifact(self, response: artifact.Artifact) -> artifact.Artifact: """Post-rpc interceptor for get_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_artifact` interceptor runs + before the `post_get_artifact_with_metadata` interceptor. """ return response + async def post_get_artifact_with_metadata( + self, + response: artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_artifact_with_metadata` + interceptor in new development instead of the `post_get_artifact` interceptor. + When both interceptors are used, this `post_get_artifact_with_metadata` interceptor runs after the + `post_get_artifact` interceptor. The (possibly modified) response returned by + `post_get_artifact` will be passed to + `post_get_artifact_with_metadata`. + """ + return response, metadata + async def pre_get_context( self, request: metadata_service.GetContextRequest, @@ -715,12 +1025,35 @@ async def pre_get_context( async def post_get_context(self, response: context.Context) -> context.Context: """Post-rpc interceptor for get_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_context` interceptor runs + before the `post_get_context_with_metadata` interceptor. """ return response + async def post_get_context_with_metadata( + self, + response: context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_context_with_metadata` + interceptor in new development instead of the `post_get_context` interceptor. + When both interceptors are used, this `post_get_context_with_metadata` interceptor runs after the + `post_get_context` interceptor. The (possibly modified) response returned by + `post_get_context` will be passed to + `post_get_context_with_metadata`. + """ + return response, metadata + async def pre_get_execution( self, request: metadata_service.GetExecutionRequest, @@ -740,12 +1073,35 @@ async def post_get_execution( ) -> execution.Execution: """Post-rpc interceptor for get_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_execution` interceptor runs + before the `post_get_execution_with_metadata` interceptor. """ return response + async def post_get_execution_with_metadata( + self, + response: execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_execution_with_metadata` + interceptor in new development instead of the `post_get_execution` interceptor. + When both interceptors are used, this `post_get_execution_with_metadata` interceptor runs after the + `post_get_execution` interceptor. The (possibly modified) response returned by + `post_get_execution` will be passed to + `post_get_execution_with_metadata`. + """ + return response, metadata + async def pre_get_metadata_schema( self, request: metadata_service.GetMetadataSchemaRequest, @@ -766,12 +1122,35 @@ async def post_get_metadata_schema( ) -> metadata_schema.MetadataSchema: """Post-rpc interceptor for get_metadata_schema - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_schema_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_schema` interceptor runs + before the `post_get_metadata_schema_with_metadata` interceptor. """ return response + async def post_get_metadata_schema_with_metadata( + self, + response: metadata_schema.MetadataSchema, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_schema.MetadataSchema, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_schema + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_schema_with_metadata` + interceptor in new development instead of the `post_get_metadata_schema` interceptor. + When both interceptors are used, this `post_get_metadata_schema_with_metadata` interceptor runs after the + `post_get_metadata_schema` interceptor. The (possibly modified) response returned by + `post_get_metadata_schema` will be passed to + `post_get_metadata_schema_with_metadata`. + """ + return response, metadata + async def pre_get_metadata_store( self, request: metadata_service.GetMetadataStoreRequest, @@ -792,12 +1171,35 @@ async def post_get_metadata_store( ) -> metadata_store.MetadataStore: """Post-rpc interceptor for get_metadata_store - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_metadata_store_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_get_metadata_store` interceptor runs + before the `post_get_metadata_store_with_metadata` interceptor. """ return response + async def post_get_metadata_store_with_metadata( + self, + response: metadata_store.MetadataStore, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[metadata_store.MetadataStore, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_metadata_store + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_get_metadata_store_with_metadata` + interceptor in new development instead of the `post_get_metadata_store` interceptor. + When both interceptors are used, this `post_get_metadata_store_with_metadata` interceptor runs after the + `post_get_metadata_store` interceptor. The (possibly modified) response returned by + `post_get_metadata_store` will be passed to + `post_get_metadata_store_with_metadata`. + """ + return response, metadata + async def pre_list_artifacts( self, request: metadata_service.ListArtifactsRequest, @@ -817,12 +1219,37 @@ async def post_list_artifacts( ) -> metadata_service.ListArtifactsResponse: """Post-rpc interceptor for list_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_artifacts` interceptor runs + before the `post_list_artifacts_with_metadata` interceptor. """ return response + async def post_list_artifacts_with_metadata( + self, + response: metadata_service.ListArtifactsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListArtifactsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_artifacts_with_metadata` + interceptor in new development instead of the `post_list_artifacts` interceptor. + When both interceptors are used, this `post_list_artifacts_with_metadata` interceptor runs after the + `post_list_artifacts` interceptor. The (possibly modified) response returned by + `post_list_artifacts` will be passed to + `post_list_artifacts_with_metadata`. + """ + return response, metadata + async def pre_list_contexts( self, request: metadata_service.ListContextsRequest, @@ -842,12 +1269,37 @@ async def post_list_contexts( ) -> metadata_service.ListContextsResponse: """Post-rpc interceptor for list_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_contexts` interceptor runs + before the `post_list_contexts_with_metadata` interceptor. """ return response + async def post_list_contexts_with_metadata( + self, + response: metadata_service.ListContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListContextsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_contexts_with_metadata` + interceptor in new development instead of the `post_list_contexts` interceptor. + When both interceptors are used, this `post_list_contexts_with_metadata` interceptor runs after the + `post_list_contexts` interceptor. The (possibly modified) response returned by + `post_list_contexts` will be passed to + `post_list_contexts_with_metadata`. + """ + return response, metadata + async def pre_list_executions( self, request: metadata_service.ListExecutionsRequest, @@ -867,12 +1319,37 @@ async def post_list_executions( ) -> metadata_service.ListExecutionsResponse: """Post-rpc interceptor for list_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_executions` interceptor runs + before the `post_list_executions_with_metadata` interceptor. """ return response + async def post_list_executions_with_metadata( + self, + response: metadata_service.ListExecutionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListExecutionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_executions_with_metadata` + interceptor in new development instead of the `post_list_executions` interceptor. + When both interceptors are used, this `post_list_executions_with_metadata` interceptor runs after the + `post_list_executions` interceptor. The (possibly modified) response returned by + `post_list_executions` will be passed to + `post_list_executions_with_metadata`. + """ + return response, metadata + async def pre_list_metadata_schemas( self, request: metadata_service.ListMetadataSchemasRequest, @@ -893,12 +1370,38 @@ async def post_list_metadata_schemas( ) -> metadata_service.ListMetadataSchemasResponse: """Post-rpc interceptor for list_metadata_schemas - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_schemas_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_schemas` interceptor runs + before the `post_list_metadata_schemas_with_metadata` interceptor. """ return response + async def post_list_metadata_schemas_with_metadata( + self, + response: metadata_service.ListMetadataSchemasResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataSchemasResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_schemas + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_schemas_with_metadata` + interceptor in new development instead of the `post_list_metadata_schemas` interceptor. + When both interceptors are used, this `post_list_metadata_schemas_with_metadata` interceptor runs after the + `post_list_metadata_schemas` interceptor. The (possibly modified) response returned by + `post_list_metadata_schemas` will be passed to + `post_list_metadata_schemas_with_metadata`. + """ + return response, metadata + async def pre_list_metadata_stores( self, request: metadata_service.ListMetadataStoresRequest, @@ -919,12 +1422,38 @@ async def post_list_metadata_stores( ) -> metadata_service.ListMetadataStoresResponse: """Post-rpc interceptor for list_metadata_stores - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_metadata_stores_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_list_metadata_stores` interceptor runs + before the `post_list_metadata_stores_with_metadata` interceptor. """ return response + async def post_list_metadata_stores_with_metadata( + self, + response: metadata_service.ListMetadataStoresResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.ListMetadataStoresResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_metadata_stores + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_list_metadata_stores_with_metadata` + interceptor in new development instead of the `post_list_metadata_stores` interceptor. + When both interceptors are used, this `post_list_metadata_stores_with_metadata` interceptor runs after the + `post_list_metadata_stores` interceptor. The (possibly modified) response returned by + `post_list_metadata_stores` will be passed to + `post_list_metadata_stores_with_metadata`. + """ + return response, metadata + async def pre_purge_artifacts( self, request: metadata_service.PurgeArtifactsRequest, @@ -944,12 +1473,35 @@ async def post_purge_artifacts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_artifacts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_artifacts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_artifacts` interceptor runs + before the `post_purge_artifacts_with_metadata` interceptor. """ return response + async def post_purge_artifacts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_artifacts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_artifacts_with_metadata` + interceptor in new development instead of the `post_purge_artifacts` interceptor. + When both interceptors are used, this `post_purge_artifacts_with_metadata` interceptor runs after the + `post_purge_artifacts` interceptor. The (possibly modified) response returned by + `post_purge_artifacts` will be passed to + `post_purge_artifacts_with_metadata`. + """ + return response, metadata + async def pre_purge_contexts( self, request: metadata_service.PurgeContextsRequest, @@ -969,12 +1521,35 @@ async def post_purge_contexts( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_contexts` interceptor runs + before the `post_purge_contexts_with_metadata` interceptor. """ return response + async def post_purge_contexts_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_contexts_with_metadata` + interceptor in new development instead of the `post_purge_contexts` interceptor. + When both interceptors are used, this `post_purge_contexts_with_metadata` interceptor runs after the + `post_purge_contexts` interceptor. The (possibly modified) response returned by + `post_purge_contexts` will be passed to + `post_purge_contexts_with_metadata`. + """ + return response, metadata + async def pre_purge_executions( self, request: metadata_service.PurgeExecutionsRequest, @@ -994,12 +1569,35 @@ async def post_purge_executions( ) -> operations_pb2.Operation: """Post-rpc interceptor for purge_executions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_purge_executions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_purge_executions` interceptor runs + before the `post_purge_executions_with_metadata` interceptor. """ return response + async def post_purge_executions_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for purge_executions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_purge_executions_with_metadata` + interceptor in new development instead of the `post_purge_executions` interceptor. + When both interceptors are used, this `post_purge_executions_with_metadata` interceptor runs after the + `post_purge_executions` interceptor. The (possibly modified) response returned by + `post_purge_executions` will be passed to + `post_purge_executions_with_metadata`. + """ + return response, metadata + async def pre_query_artifact_lineage_subgraph( self, request: metadata_service.QueryArtifactLineageSubgraphRequest, @@ -1020,12 +1618,37 @@ async def post_query_artifact_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_artifact_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_artifact_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_artifact_lineage_subgraph` interceptor runs + before the `post_query_artifact_lineage_subgraph_with_metadata` interceptor. """ return response + async def post_query_artifact_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_artifact_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_artifact_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_artifact_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_artifact_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_artifact_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_artifact_lineage_subgraph` will be passed to + `post_query_artifact_lineage_subgraph_with_metadata`. + """ + return response, metadata + async def pre_query_context_lineage_subgraph( self, request: metadata_service.QueryContextLineageSubgraphRequest, @@ -1046,12 +1669,37 @@ async def post_query_context_lineage_subgraph( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_context_lineage_subgraph - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_context_lineage_subgraph_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_context_lineage_subgraph` interceptor runs + before the `post_query_context_lineage_subgraph_with_metadata` interceptor. """ return response + async def post_query_context_lineage_subgraph_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_context_lineage_subgraph + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_context_lineage_subgraph_with_metadata` + interceptor in new development instead of the `post_query_context_lineage_subgraph` interceptor. + When both interceptors are used, this `post_query_context_lineage_subgraph_with_metadata` interceptor runs after the + `post_query_context_lineage_subgraph` interceptor. The (possibly modified) response returned by + `post_query_context_lineage_subgraph` will be passed to + `post_query_context_lineage_subgraph_with_metadata`. + """ + return response, metadata + async def pre_query_execution_inputs_and_outputs( self, request: metadata_service.QueryExecutionInputsAndOutputsRequest, @@ -1072,12 +1720,37 @@ async def post_query_execution_inputs_and_outputs( ) -> lineage_subgraph.LineageSubgraph: """Post-rpc interceptor for query_execution_inputs_and_outputs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_execution_inputs_and_outputs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_query_execution_inputs_and_outputs` interceptor runs + before the `post_query_execution_inputs_and_outputs_with_metadata` interceptor. """ return response + async def post_query_execution_inputs_and_outputs_with_metadata( + self, + response: lineage_subgraph.LineageSubgraph, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + lineage_subgraph.LineageSubgraph, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for query_execution_inputs_and_outputs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_query_execution_inputs_and_outputs_with_metadata` + interceptor in new development instead of the `post_query_execution_inputs_and_outputs` interceptor. + When both interceptors are used, this `post_query_execution_inputs_and_outputs_with_metadata` interceptor runs after the + `post_query_execution_inputs_and_outputs` interceptor. The (possibly modified) response returned by + `post_query_execution_inputs_and_outputs` will be passed to + `post_query_execution_inputs_and_outputs_with_metadata`. + """ + return response, metadata + async def pre_remove_context_children( self, request: metadata_service.RemoveContextChildrenRequest, @@ -1098,12 +1771,38 @@ async def post_remove_context_children( ) -> metadata_service.RemoveContextChildrenResponse: """Post-rpc interceptor for remove_context_children - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_remove_context_children_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_remove_context_children` interceptor runs + before the `post_remove_context_children_with_metadata` interceptor. """ return response + async def post_remove_context_children_with_metadata( + self, + response: metadata_service.RemoveContextChildrenResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + metadata_service.RemoveContextChildrenResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for remove_context_children + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_remove_context_children_with_metadata` + interceptor in new development instead of the `post_remove_context_children` interceptor. + When both interceptors are used, this `post_remove_context_children_with_metadata` interceptor runs after the + `post_remove_context_children` interceptor. The (possibly modified) response returned by + `post_remove_context_children` will be passed to + `post_remove_context_children_with_metadata`. + """ + return response, metadata + async def pre_update_artifact( self, request: metadata_service.UpdateArtifactRequest, @@ -1123,12 +1822,35 @@ async def post_update_artifact( ) -> gca_artifact.Artifact: """Post-rpc interceptor for update_artifact - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_artifact_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_artifact` interceptor runs + before the `post_update_artifact_with_metadata` interceptor. """ return response + async def post_update_artifact_with_metadata( + self, + response: gca_artifact.Artifact, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_artifact.Artifact, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_artifact + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_artifact_with_metadata` + interceptor in new development instead of the `post_update_artifact` interceptor. + When both interceptors are used, this `post_update_artifact_with_metadata` interceptor runs after the + `post_update_artifact` interceptor. The (possibly modified) response returned by + `post_update_artifact` will be passed to + `post_update_artifact_with_metadata`. + """ + return response, metadata + async def pre_update_context( self, request: metadata_service.UpdateContextRequest, @@ -1148,12 +1870,35 @@ async def post_update_context( ) -> gca_context.Context: """Post-rpc interceptor for update_context - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_context_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_context` interceptor runs + before the `post_update_context_with_metadata` interceptor. """ return response + async def post_update_context_with_metadata( + self, + response: gca_context.Context, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_context.Context, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_context + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_context_with_metadata` + interceptor in new development instead of the `post_update_context` interceptor. + When both interceptors are used, this `post_update_context_with_metadata` interceptor runs after the + `post_update_context` interceptor. The (possibly modified) response returned by + `post_update_context` will be passed to + `post_update_context_with_metadata`. + """ + return response, metadata + async def pre_update_execution( self, request: metadata_service.UpdateExecutionRequest, @@ -1173,12 +1918,35 @@ async def post_update_execution( ) -> gca_execution.Execution: """Post-rpc interceptor for update_execution - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_execution_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MetadataService server but before - it is returned to user code. + it is returned to user code. This `post_update_execution` interceptor runs + before the `post_update_execution_with_metadata` interceptor. """ return response + async def post_update_execution_with_metadata( + self, + response: gca_execution.Execution, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_execution.Execution, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_execution + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MetadataService server but before it is returned to user code. + + We recommend only using this `post_update_execution_with_metadata` + interceptor in new development instead of the `post_update_execution` interceptor. + When both interceptors are used, this `post_update_execution_with_metadata` interceptor runs after the + `post_update_execution` interceptor. The (possibly modified) response returned by + `post_update_execution` will be passed to + `post_update_execution_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1854,6 +2622,13 @@ async def __call__( resp = await self._interceptor.post_add_context_artifacts_and_executions( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_add_context_artifacts_and_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2015,6 +2790,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2176,6 +2955,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_execution_events(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_execution_events_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2337,6 +3120,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2496,6 +3283,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2655,6 +3446,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2812,6 +3607,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2974,6 +3773,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3130,6 +3933,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3286,6 +4093,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3442,6 +4253,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3596,6 +4411,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3748,6 +4567,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3900,6 +4723,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4053,6 +4880,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4204,6 +5035,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_metadata_schema(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_metadata_schema_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4360,6 +5195,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_metadata_store(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_metadata_store_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4515,6 +5354,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4672,6 +5515,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4829,6 +5676,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4984,6 +5835,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_metadata_schemas(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_metadata_schemas_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5139,6 +5994,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_metadata_stores(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_metadata_stores_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5303,6 +6162,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_artifacts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_artifacts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5465,6 +6328,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5627,6 +6494,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_purge_executions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_purge_executions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5787,6 +6658,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_artifact_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_artifact_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5947,6 +6825,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_context_lineage_subgraph(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_context_lineage_subgraph_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6109,6 +6994,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_execution_inputs_and_outputs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_query_execution_inputs_and_outputs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6270,6 +7162,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_remove_context_children(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_remove_context_children_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6431,6 +7330,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_artifact(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_artifact_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6590,6 +7493,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_context(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_context_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6749,6 +7656,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_execution(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_execution_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/migration_service/client.py b/google/cloud/aiplatform_v1beta1/services/migration_service/client.py index 25faa0390b..ebbb6c7208 100644 --- a/google/cloud/aiplatform_v1beta1/services/migration_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/migration_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -633,6 +635,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1189,16 +1218,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1244,16 +1277,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1416,16 +1453,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1537,16 +1578,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1659,16 +1704,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1719,16 +1768,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1774,16 +1827,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1829,16 +1886,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest.py index ad9cc51149..527625a180 100644 --- a/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest.py @@ -121,12 +121,35 @@ def post_batch_migrate_resources( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_migrate_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_migrate_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_batch_migrate_resources` interceptor runs + before the `post_batch_migrate_resources_with_metadata` interceptor. """ return response + def post_batch_migrate_resources_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_migrate_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_batch_migrate_resources_with_metadata` + interceptor in new development instead of the `post_batch_migrate_resources` interceptor. + When both interceptors are used, this `post_batch_migrate_resources_with_metadata` interceptor runs after the + `post_batch_migrate_resources` interceptor. The (possibly modified) response returned by + `post_batch_migrate_resources` will be passed to + `post_batch_migrate_resources_with_metadata`. + """ + return response, metadata + def pre_search_migratable_resources( self, request: migration_service.SearchMigratableResourcesRequest, @@ -147,12 +170,38 @@ def post_search_migratable_resources( ) -> migration_service.SearchMigratableResourcesResponse: """Post-rpc interceptor for search_migratable_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_migratable_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_search_migratable_resources` interceptor runs + before the `post_search_migratable_resources_with_metadata` interceptor. """ return response + def post_search_migratable_resources_with_metadata( + self, + response: migration_service.SearchMigratableResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + migration_service.SearchMigratableResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_migratable_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_search_migratable_resources_with_metadata` + interceptor in new development instead of the `post_search_migratable_resources` interceptor. + When both interceptors are used, this `post_search_migratable_resources_with_metadata` interceptor runs after the + `post_search_migratable_resources` interceptor. The (possibly modified) response returned by + `post_search_migratable_resources` will be passed to + `post_search_migratable_resources_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2702,6 +2751,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_migrate_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_migrate_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2858,6 +2911,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_migratable_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_search_migratable_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest_asyncio.py index b621a19f2e..9b592601a5 100644 --- a/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/migration_service/transports/rest_asyncio.py @@ -139,12 +139,35 @@ async def post_batch_migrate_resources( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_migrate_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_migrate_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_batch_migrate_resources` interceptor runs + before the `post_batch_migrate_resources_with_metadata` interceptor. """ return response + async def post_batch_migrate_resources_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_migrate_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_batch_migrate_resources_with_metadata` + interceptor in new development instead of the `post_batch_migrate_resources` interceptor. + When both interceptors are used, this `post_batch_migrate_resources_with_metadata` interceptor runs after the + `post_batch_migrate_resources` interceptor. The (possibly modified) response returned by + `post_batch_migrate_resources` will be passed to + `post_batch_migrate_resources_with_metadata`. + """ + return response, metadata + async def pre_search_migratable_resources( self, request: migration_service.SearchMigratableResourcesRequest, @@ -165,12 +188,38 @@ async def post_search_migratable_resources( ) -> migration_service.SearchMigratableResourcesResponse: """Post-rpc interceptor for search_migratable_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_migratable_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the MigrationService server but before - it is returned to user code. + it is returned to user code. This `post_search_migratable_resources` interceptor runs + before the `post_search_migratable_resources_with_metadata` interceptor. """ return response + async def post_search_migratable_resources_with_metadata( + self, + response: migration_service.SearchMigratableResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + migration_service.SearchMigratableResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_migratable_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the MigrationService server but before it is returned to user code. + + We recommend only using this `post_search_migratable_resources_with_metadata` + interceptor in new development instead of the `post_search_migratable_resources` interceptor. + When both interceptors are used, this `post_search_migratable_resources_with_metadata` interceptor runs after the + `post_search_migratable_resources` interceptor. The (possibly modified) response returned by + `post_search_migratable_resources` will be passed to + `post_search_migratable_resources_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -690,6 +739,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_migrate_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_migrate_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -850,6 +906,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_migratable_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_migratable_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/model_garden_service/client.py b/google/cloud/aiplatform_v1beta1/services/model_garden_service/client.py index 0276130295..7539a58c9e 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_garden_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/model_garden_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -566,6 +568,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1180,16 +1209,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1235,16 +1268,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1407,16 +1444,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1528,16 +1569,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1650,16 +1695,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1710,16 +1759,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1765,16 +1818,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1820,16 +1877,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest.py index a345991c15..26d3d177e9 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest.py @@ -130,12 +130,35 @@ def post_deploy_publisher_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_publisher_model` interceptor runs + before the `post_deploy_publisher_model_with_metadata` interceptor. """ return response + def post_deploy_publisher_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_deploy_publisher_model_with_metadata` + interceptor in new development instead of the `post_deploy_publisher_model` interceptor. + When both interceptors are used, this `post_deploy_publisher_model_with_metadata` interceptor runs after the + `post_deploy_publisher_model` interceptor. The (possibly modified) response returned by + `post_deploy_publisher_model` will be passed to + `post_deploy_publisher_model_with_metadata`. + """ + return response, metadata + def pre_get_publisher_model( self, request: model_garden_service.GetPublisherModelRequest, @@ -156,12 +179,35 @@ def post_get_publisher_model( ) -> publisher_model.PublisherModel: """Post-rpc interceptor for get_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_get_publisher_model` interceptor runs + before the `post_get_publisher_model_with_metadata` interceptor. """ return response + def post_get_publisher_model_with_metadata( + self, + response: publisher_model.PublisherModel, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[publisher_model.PublisherModel, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_get_publisher_model_with_metadata` + interceptor in new development instead of the `post_get_publisher_model` interceptor. + When both interceptors are used, this `post_get_publisher_model_with_metadata` interceptor runs after the + `post_get_publisher_model` interceptor. The (possibly modified) response returned by + `post_get_publisher_model` will be passed to + `post_get_publisher_model_with_metadata`. + """ + return response, metadata + def pre_list_publisher_models( self, request: model_garden_service.ListPublisherModelsRequest, @@ -182,12 +228,38 @@ def post_list_publisher_models( ) -> model_garden_service.ListPublisherModelsResponse: """Post-rpc interceptor for list_publisher_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_publisher_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_list_publisher_models` interceptor runs + before the `post_list_publisher_models_with_metadata` interceptor. """ return response + def post_list_publisher_models_with_metadata( + self, + response: model_garden_service.ListPublisherModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_garden_service.ListPublisherModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_publisher_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_list_publisher_models_with_metadata` + interceptor in new development instead of the `post_list_publisher_models` interceptor. + When both interceptors are used, this `post_list_publisher_models_with_metadata` interceptor runs after the + `post_list_publisher_models` interceptor. The (possibly modified) response returned by + `post_list_publisher_models` will be passed to + `post_list_publisher_models_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2736,6 +2808,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_deploy_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_deploy_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2881,6 +2957,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3030,6 +3110,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_publisher_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_publisher_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest_asyncio.py index f7b11aff9b..f5db69482b 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/model_garden_service/transports/rest_asyncio.py @@ -148,12 +148,35 @@ async def post_deploy_publisher_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for deploy_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_deploy_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_deploy_publisher_model` interceptor runs + before the `post_deploy_publisher_model_with_metadata` interceptor. """ return response + async def post_deploy_publisher_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for deploy_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_deploy_publisher_model_with_metadata` + interceptor in new development instead of the `post_deploy_publisher_model` interceptor. + When both interceptors are used, this `post_deploy_publisher_model_with_metadata` interceptor runs after the + `post_deploy_publisher_model` interceptor. The (possibly modified) response returned by + `post_deploy_publisher_model` will be passed to + `post_deploy_publisher_model_with_metadata`. + """ + return response, metadata + async def pre_get_publisher_model( self, request: model_garden_service.GetPublisherModelRequest, @@ -174,12 +197,35 @@ async def post_get_publisher_model( ) -> publisher_model.PublisherModel: """Post-rpc interceptor for get_publisher_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_publisher_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_get_publisher_model` interceptor runs + before the `post_get_publisher_model_with_metadata` interceptor. """ return response + async def post_get_publisher_model_with_metadata( + self, + response: publisher_model.PublisherModel, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[publisher_model.PublisherModel, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_publisher_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_get_publisher_model_with_metadata` + interceptor in new development instead of the `post_get_publisher_model` interceptor. + When both interceptors are used, this `post_get_publisher_model_with_metadata` interceptor runs after the + `post_get_publisher_model` interceptor. The (possibly modified) response returned by + `post_get_publisher_model` will be passed to + `post_get_publisher_model_with_metadata`. + """ + return response, metadata + async def pre_list_publisher_models( self, request: model_garden_service.ListPublisherModelsRequest, @@ -200,12 +246,38 @@ async def post_list_publisher_models( ) -> model_garden_service.ListPublisherModelsResponse: """Post-rpc interceptor for list_publisher_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_publisher_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelGardenService server but before - it is returned to user code. + it is returned to user code. This `post_list_publisher_models` interceptor runs + before the `post_list_publisher_models_with_metadata` interceptor. """ return response + async def post_list_publisher_models_with_metadata( + self, + response: model_garden_service.ListPublisherModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_garden_service.ListPublisherModelsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_publisher_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelGardenService server but before it is returned to user code. + + We recommend only using this `post_list_publisher_models_with_metadata` + interceptor in new development instead of the `post_list_publisher_models` interceptor. + When both interceptors are used, this `post_list_publisher_models_with_metadata` interceptor runs after the + `post_list_publisher_models` interceptor. The (possibly modified) response returned by + `post_list_publisher_models` will be passed to + `post_list_publisher_models_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -729,6 +801,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_deploy_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_deploy_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -880,6 +956,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_publisher_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_publisher_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1033,6 +1113,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_publisher_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_publisher_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/client.py b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/client.py index 0fb844712d..3e2a708632 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -679,6 +681,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2355,16 +2384,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2410,16 +2443,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2582,16 +2619,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2703,16 +2744,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2825,16 +2870,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2885,16 +2934,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2940,16 +2993,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2995,16 +3052,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest.py index 87835d8c18..ab70f057c2 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest.py @@ -198,12 +198,35 @@ def post_create_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_monitor` interceptor runs + before the `post_create_model_monitor_with_metadata` interceptor. """ return response + def post_create_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_create_model_monitor_with_metadata` + interceptor in new development instead of the `post_create_model_monitor` interceptor. + When both interceptors are used, this `post_create_model_monitor_with_metadata` interceptor runs after the + `post_create_model_monitor` interceptor. The (possibly modified) response returned by + `post_create_model_monitor` will be passed to + `post_create_model_monitor_with_metadata`. + """ + return response, metadata + def pre_create_model_monitoring_job( self, request: model_monitoring_service.CreateModelMonitoringJobRequest, @@ -224,12 +247,38 @@ def post_create_model_monitoring_job( ) -> gca_model_monitoring_job.ModelMonitoringJob: """Post-rpc interceptor for create_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_monitoring_job` interceptor runs + before the `post_create_model_monitoring_job_with_metadata` interceptor. """ return response + def post_create_model_monitoring_job_with_metadata( + self, + response: gca_model_monitoring_job.ModelMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_monitoring_job.ModelMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_create_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_monitoring_job` will be passed to + `post_create_model_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_delete_model_monitor( self, request: model_monitoring_service.DeleteModelMonitorRequest, @@ -250,12 +299,35 @@ def post_delete_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_monitor` interceptor runs + before the `post_delete_model_monitor_with_metadata` interceptor. """ return response + def post_delete_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_delete_model_monitor_with_metadata` + interceptor in new development instead of the `post_delete_model_monitor` interceptor. + When both interceptors are used, this `post_delete_model_monitor_with_metadata` interceptor runs after the + `post_delete_model_monitor` interceptor. The (possibly modified) response returned by + `post_delete_model_monitor` will be passed to + `post_delete_model_monitor_with_metadata`. + """ + return response, metadata + def pre_delete_model_monitoring_job( self, request: model_monitoring_service.DeleteModelMonitoringJobRequest, @@ -276,12 +348,35 @@ def post_delete_model_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_monitoring_job` interceptor runs + before the `post_delete_model_monitoring_job_with_metadata` interceptor. """ return response + def post_delete_model_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_delete_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_monitoring_job` will be passed to + `post_delete_model_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_get_model_monitor( self, request: model_monitoring_service.GetModelMonitorRequest, @@ -302,12 +397,35 @@ def post_get_model_monitor( ) -> model_monitor.ModelMonitor: """Post-rpc interceptor for get_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_monitor` interceptor runs + before the `post_get_model_monitor_with_metadata` interceptor. """ return response + def post_get_model_monitor_with_metadata( + self, + response: model_monitor.ModelMonitor, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[model_monitor.ModelMonitor, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_get_model_monitor_with_metadata` + interceptor in new development instead of the `post_get_model_monitor` interceptor. + When both interceptors are used, this `post_get_model_monitor_with_metadata` interceptor runs after the + `post_get_model_monitor` interceptor. The (possibly modified) response returned by + `post_get_model_monitor` will be passed to + `post_get_model_monitor_with_metadata`. + """ + return response, metadata + def pre_get_model_monitoring_job( self, request: model_monitoring_service.GetModelMonitoringJobRequest, @@ -328,12 +446,37 @@ def post_get_model_monitoring_job( ) -> model_monitoring_job.ModelMonitoringJob: """Post-rpc interceptor for get_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_monitoring_job` interceptor runs + before the `post_get_model_monitoring_job_with_metadata` interceptor. """ return response + def post_get_model_monitoring_job_with_metadata( + self, + response: model_monitoring_job.ModelMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_job.ModelMonitoringJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_get_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_monitoring_job` will be passed to + `post_get_model_monitoring_job_with_metadata`. + """ + return response, metadata + def pre_list_model_monitoring_jobs( self, request: model_monitoring_service.ListModelMonitoringJobsRequest, @@ -354,12 +497,38 @@ def post_list_model_monitoring_jobs( ) -> model_monitoring_service.ListModelMonitoringJobsResponse: """Post-rpc interceptor for list_model_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_monitoring_jobs` interceptor runs + before the `post_list_model_monitoring_jobs_with_metadata` interceptor. """ return response + def post_list_model_monitoring_jobs_with_metadata( + self, + response: model_monitoring_service.ListModelMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.ListModelMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_list_model_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_monitoring_jobs` will be passed to + `post_list_model_monitoring_jobs_with_metadata`. + """ + return response, metadata + def pre_list_model_monitors( self, request: model_monitoring_service.ListModelMonitorsRequest, @@ -380,12 +549,38 @@ def post_list_model_monitors( ) -> model_monitoring_service.ListModelMonitorsResponse: """Post-rpc interceptor for list_model_monitors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_monitors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_monitors` interceptor runs + before the `post_list_model_monitors_with_metadata` interceptor. """ return response + def post_list_model_monitors_with_metadata( + self, + response: model_monitoring_service.ListModelMonitorsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.ListModelMonitorsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_monitors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_list_model_monitors_with_metadata` + interceptor in new development instead of the `post_list_model_monitors` interceptor. + When both interceptors are used, this `post_list_model_monitors_with_metadata` interceptor runs after the + `post_list_model_monitors` interceptor. The (possibly modified) response returned by + `post_list_model_monitors` will be passed to + `post_list_model_monitors_with_metadata`. + """ + return response, metadata + def pre_search_model_monitoring_alerts( self, request: model_monitoring_service.SearchModelMonitoringAlertsRequest, @@ -406,12 +601,38 @@ def post_search_model_monitoring_alerts( ) -> model_monitoring_service.SearchModelMonitoringAlertsResponse: """Post-rpc interceptor for search_model_monitoring_alerts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_monitoring_alerts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_monitoring_alerts` interceptor runs + before the `post_search_model_monitoring_alerts_with_metadata` interceptor. """ return response + def post_search_model_monitoring_alerts_with_metadata( + self, + response: model_monitoring_service.SearchModelMonitoringAlertsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.SearchModelMonitoringAlertsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_monitoring_alerts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_search_model_monitoring_alerts_with_metadata` + interceptor in new development instead of the `post_search_model_monitoring_alerts` interceptor. + When both interceptors are used, this `post_search_model_monitoring_alerts_with_metadata` interceptor runs after the + `post_search_model_monitoring_alerts` interceptor. The (possibly modified) response returned by + `post_search_model_monitoring_alerts` will be passed to + `post_search_model_monitoring_alerts_with_metadata`. + """ + return response, metadata + def pre_search_model_monitoring_stats( self, request: model_monitoring_service.SearchModelMonitoringStatsRequest, @@ -432,12 +653,38 @@ def post_search_model_monitoring_stats( ) -> model_monitoring_service.SearchModelMonitoringStatsResponse: """Post-rpc interceptor for search_model_monitoring_stats - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_monitoring_stats_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_monitoring_stats` interceptor runs + before the `post_search_model_monitoring_stats_with_metadata` interceptor. """ return response + def post_search_model_monitoring_stats_with_metadata( + self, + response: model_monitoring_service.SearchModelMonitoringStatsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.SearchModelMonitoringStatsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_monitoring_stats + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_search_model_monitoring_stats_with_metadata` + interceptor in new development instead of the `post_search_model_monitoring_stats` interceptor. + When both interceptors are used, this `post_search_model_monitoring_stats_with_metadata` interceptor runs after the + `post_search_model_monitoring_stats` interceptor. The (possibly modified) response returned by + `post_search_model_monitoring_stats` will be passed to + `post_search_model_monitoring_stats_with_metadata`. + """ + return response, metadata + def pre_update_model_monitor( self, request: model_monitoring_service.UpdateModelMonitorRequest, @@ -458,12 +705,35 @@ def post_update_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_monitor` interceptor runs + before the `post_update_model_monitor_with_metadata` interceptor. """ return response + def post_update_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_update_model_monitor_with_metadata` + interceptor in new development instead of the `post_update_model_monitor` interceptor. + When both interceptors are used, this `post_update_model_monitor_with_metadata` interceptor runs after the + `post_update_model_monitor` interceptor. The (possibly modified) response returned by + `post_update_model_monitor` will be passed to + `post_update_model_monitor_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3014,6 +3284,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3169,6 +3443,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3319,6 +3597,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3466,6 +3748,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3620,6 +3906,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3768,6 +4058,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3918,6 +4212,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_monitoring_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4069,6 +4367,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_monitors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_monitors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4231,6 +4533,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_model_monitoring_alerts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_search_model_monitoring_alerts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4391,6 +4700,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_search_model_monitoring_stats(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_search_model_monitoring_stats_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4547,6 +4863,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest_asyncio.py index 8063744467..b0803934eb 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/model_monitoring_service/transports/rest_asyncio.py @@ -216,12 +216,35 @@ async def post_create_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_monitor` interceptor runs + before the `post_create_model_monitor_with_metadata` interceptor. """ return response + async def post_create_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_create_model_monitor_with_metadata` + interceptor in new development instead of the `post_create_model_monitor` interceptor. + When both interceptors are used, this `post_create_model_monitor_with_metadata` interceptor runs after the + `post_create_model_monitor` interceptor. The (possibly modified) response returned by + `post_create_model_monitor` will be passed to + `post_create_model_monitor_with_metadata`. + """ + return response, metadata + async def pre_create_model_monitoring_job( self, request: model_monitoring_service.CreateModelMonitoringJobRequest, @@ -242,12 +265,38 @@ async def post_create_model_monitoring_job( ) -> gca_model_monitoring_job.ModelMonitoringJob: """Post-rpc interceptor for create_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_create_model_monitoring_job` interceptor runs + before the `post_create_model_monitoring_job_with_metadata` interceptor. """ return response + async def post_create_model_monitoring_job_with_metadata( + self, + response: gca_model_monitoring_job.ModelMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_monitoring_job.ModelMonitoringJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_create_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_create_model_monitoring_job` interceptor. + When both interceptors are used, this `post_create_model_monitoring_job_with_metadata` interceptor runs after the + `post_create_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_create_model_monitoring_job` will be passed to + `post_create_model_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_delete_model_monitor( self, request: model_monitoring_service.DeleteModelMonitorRequest, @@ -268,12 +317,35 @@ async def post_delete_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_monitor` interceptor runs + before the `post_delete_model_monitor_with_metadata` interceptor. """ return response + async def post_delete_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_delete_model_monitor_with_metadata` + interceptor in new development instead of the `post_delete_model_monitor` interceptor. + When both interceptors are used, this `post_delete_model_monitor_with_metadata` interceptor runs after the + `post_delete_model_monitor` interceptor. The (possibly modified) response returned by + `post_delete_model_monitor` will be passed to + `post_delete_model_monitor_with_metadata`. + """ + return response, metadata + async def pre_delete_model_monitoring_job( self, request: model_monitoring_service.DeleteModelMonitoringJobRequest, @@ -294,12 +366,35 @@ async def post_delete_model_monitoring_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_monitoring_job` interceptor runs + before the `post_delete_model_monitoring_job_with_metadata` interceptor. """ return response + async def post_delete_model_monitoring_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_delete_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_delete_model_monitoring_job` interceptor. + When both interceptors are used, this `post_delete_model_monitoring_job_with_metadata` interceptor runs after the + `post_delete_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_delete_model_monitoring_job` will be passed to + `post_delete_model_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_get_model_monitor( self, request: model_monitoring_service.GetModelMonitorRequest, @@ -320,12 +415,35 @@ async def post_get_model_monitor( ) -> model_monitor.ModelMonitor: """Post-rpc interceptor for get_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_monitor` interceptor runs + before the `post_get_model_monitor_with_metadata` interceptor. """ return response + async def post_get_model_monitor_with_metadata( + self, + response: model_monitor.ModelMonitor, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[model_monitor.ModelMonitor, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_get_model_monitor_with_metadata` + interceptor in new development instead of the `post_get_model_monitor` interceptor. + When both interceptors are used, this `post_get_model_monitor_with_metadata` interceptor runs after the + `post_get_model_monitor` interceptor. The (possibly modified) response returned by + `post_get_model_monitor` will be passed to + `post_get_model_monitor_with_metadata`. + """ + return response, metadata + async def pre_get_model_monitoring_job( self, request: model_monitoring_service.GetModelMonitoringJobRequest, @@ -346,12 +464,37 @@ async def post_get_model_monitoring_job( ) -> model_monitoring_job.ModelMonitoringJob: """Post-rpc interceptor for get_model_monitoring_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_monitoring_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_monitoring_job` interceptor runs + before the `post_get_model_monitoring_job_with_metadata` interceptor. """ return response + async def post_get_model_monitoring_job_with_metadata( + self, + response: model_monitoring_job.ModelMonitoringJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_job.ModelMonitoringJob, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_monitoring_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_get_model_monitoring_job_with_metadata` + interceptor in new development instead of the `post_get_model_monitoring_job` interceptor. + When both interceptors are used, this `post_get_model_monitoring_job_with_metadata` interceptor runs after the + `post_get_model_monitoring_job` interceptor. The (possibly modified) response returned by + `post_get_model_monitoring_job` will be passed to + `post_get_model_monitoring_job_with_metadata`. + """ + return response, metadata + async def pre_list_model_monitoring_jobs( self, request: model_monitoring_service.ListModelMonitoringJobsRequest, @@ -372,12 +515,38 @@ async def post_list_model_monitoring_jobs( ) -> model_monitoring_service.ListModelMonitoringJobsResponse: """Post-rpc interceptor for list_model_monitoring_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_monitoring_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_monitoring_jobs` interceptor runs + before the `post_list_model_monitoring_jobs_with_metadata` interceptor. """ return response + async def post_list_model_monitoring_jobs_with_metadata( + self, + response: model_monitoring_service.ListModelMonitoringJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.ListModelMonitoringJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_monitoring_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_list_model_monitoring_jobs_with_metadata` + interceptor in new development instead of the `post_list_model_monitoring_jobs` interceptor. + When both interceptors are used, this `post_list_model_monitoring_jobs_with_metadata` interceptor runs after the + `post_list_model_monitoring_jobs` interceptor. The (possibly modified) response returned by + `post_list_model_monitoring_jobs` will be passed to + `post_list_model_monitoring_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_model_monitors( self, request: model_monitoring_service.ListModelMonitorsRequest, @@ -398,12 +567,38 @@ async def post_list_model_monitors( ) -> model_monitoring_service.ListModelMonitorsResponse: """Post-rpc interceptor for list_model_monitors - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_monitors_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_monitors` interceptor runs + before the `post_list_model_monitors_with_metadata` interceptor. """ return response + async def post_list_model_monitors_with_metadata( + self, + response: model_monitoring_service.ListModelMonitorsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.ListModelMonitorsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_monitors + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_list_model_monitors_with_metadata` + interceptor in new development instead of the `post_list_model_monitors` interceptor. + When both interceptors are used, this `post_list_model_monitors_with_metadata` interceptor runs after the + `post_list_model_monitors` interceptor. The (possibly modified) response returned by + `post_list_model_monitors` will be passed to + `post_list_model_monitors_with_metadata`. + """ + return response, metadata + async def pre_search_model_monitoring_alerts( self, request: model_monitoring_service.SearchModelMonitoringAlertsRequest, @@ -424,12 +619,38 @@ async def post_search_model_monitoring_alerts( ) -> model_monitoring_service.SearchModelMonitoringAlertsResponse: """Post-rpc interceptor for search_model_monitoring_alerts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_monitoring_alerts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_monitoring_alerts` interceptor runs + before the `post_search_model_monitoring_alerts_with_metadata` interceptor. """ return response + async def post_search_model_monitoring_alerts_with_metadata( + self, + response: model_monitoring_service.SearchModelMonitoringAlertsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.SearchModelMonitoringAlertsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_monitoring_alerts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_search_model_monitoring_alerts_with_metadata` + interceptor in new development instead of the `post_search_model_monitoring_alerts` interceptor. + When both interceptors are used, this `post_search_model_monitoring_alerts_with_metadata` interceptor runs after the + `post_search_model_monitoring_alerts` interceptor. The (possibly modified) response returned by + `post_search_model_monitoring_alerts` will be passed to + `post_search_model_monitoring_alerts_with_metadata`. + """ + return response, metadata + async def pre_search_model_monitoring_stats( self, request: model_monitoring_service.SearchModelMonitoringStatsRequest, @@ -450,12 +671,38 @@ async def post_search_model_monitoring_stats( ) -> model_monitoring_service.SearchModelMonitoringStatsResponse: """Post-rpc interceptor for search_model_monitoring_stats - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_search_model_monitoring_stats_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_search_model_monitoring_stats` interceptor runs + before the `post_search_model_monitoring_stats_with_metadata` interceptor. """ return response + async def post_search_model_monitoring_stats_with_metadata( + self, + response: model_monitoring_service.SearchModelMonitoringStatsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_monitoring_service.SearchModelMonitoringStatsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for search_model_monitoring_stats + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_search_model_monitoring_stats_with_metadata` + interceptor in new development instead of the `post_search_model_monitoring_stats` interceptor. + When both interceptors are used, this `post_search_model_monitoring_stats_with_metadata` interceptor runs after the + `post_search_model_monitoring_stats` interceptor. The (possibly modified) response returned by + `post_search_model_monitoring_stats` will be passed to + `post_search_model_monitoring_stats_with_metadata`. + """ + return response, metadata + async def pre_update_model_monitor( self, request: model_monitoring_service.UpdateModelMonitorRequest, @@ -476,12 +723,35 @@ async def post_update_model_monitor( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_model_monitor - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_monitor_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelMonitoringService server but before - it is returned to user code. + it is returned to user code. This `post_update_model_monitor` interceptor runs + before the `post_update_model_monitor_with_metadata` interceptor. """ return response + async def post_update_model_monitor_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model_monitor + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelMonitoringService server but before it is returned to user code. + + We recommend only using this `post_update_model_monitor_with_metadata` + interceptor in new development instead of the `post_update_model_monitor` interceptor. + When both interceptors are used, this `post_update_model_monitor_with_metadata` interceptor runs after the + `post_update_model_monitor` interceptor. The (possibly modified) response returned by + `post_update_model_monitor` will be passed to + `post_update_model_monitor_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1049,6 +1319,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1212,6 +1486,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1368,6 +1649,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1525,6 +1810,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1683,6 +1975,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1839,6 +2135,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_monitoring_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_model_monitoring_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1997,6 +2300,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_monitoring_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_monitoring_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2152,6 +2462,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_monitors(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_model_monitors_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2323,6 +2637,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_model_monitoring_alerts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_model_monitoring_alerts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2492,6 +2813,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_search_model_monitoring_stats(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_search_model_monitoring_stats_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2654,6 +2982,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_model_monitor(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_model_monitor_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/async_client.py b/google/cloud/aiplatform_v1beta1/services/model_service/async_client.py index 266c7068af..d96cd7b437 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/async_client.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/async_client.py @@ -820,6 +820,138 @@ async def sample_list_model_versions(): # Done; return the response. return response + async def list_model_version_checkpoints( + self, + request: Optional[ + Union[model_service.ListModelVersionCheckpointsRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> pagers.ListModelVersionCheckpointsAsyncPager: + r"""Lists checkpoints of the specified model version. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1beta1 + + async def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1beta1.ModelServiceAsyncClient() + + # Initialize request argument(s) + request = aiplatform_v1beta1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + async for response in page_result: + print(response) + + Args: + request (Optional[Union[google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest, dict]]): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]. + name (:class:`str`): + Required. The name of the model version to list + checkpoints for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest + version will be used. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.cloud.aiplatform_v1beta1.services.model_service.pagers.ListModelVersionCheckpointsAsyncPager: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + + Iterating over this object will yield results and + resolve additional pages automatically. + + """ + # Create or coerce a protobuf request object. + # - Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, model_service.ListModelVersionCheckpointsRequest): + request = model_service.ListModelVersionCheckpointsRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._client._transport._wrapped_methods[ + self._client._transport.list_model_version_checkpoints + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # This method is paged; wrap the response in a pager, which provides + # an `__aiter__` convenience method. + response = pagers.ListModelVersionCheckpointsAsyncPager( + method=rpc, + request=request, + response=response, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + async def update_model( self, request: Optional[Union[model_service.UpdateModelRequest, dict]] = None, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/client.py b/google/cloud/aiplatform_v1beta1/services/model_service/client.py index c7d0c8a40c..976e6bc647 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -611,6 +613,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1317,6 +1346,137 @@ def sample_list_model_versions(): # Done; return the response. return response + def list_model_version_checkpoints( + self, + request: Optional[ + Union[model_service.ListModelVersionCheckpointsRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> pagers.ListModelVersionCheckpointsPager: + r"""Lists checkpoints of the specified model version. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import aiplatform_v1beta1 + + def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1beta1.ModelServiceClient() + + # Initialize request argument(s) + request = aiplatform_v1beta1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + for response in page_result: + print(response) + + Args: + request (Union[google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest, dict]): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]. + name (str): + Required. The name of the model version to list + checkpoints for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest + version will be used. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + google.cloud.aiplatform_v1beta1.services.model_service.pagers.ListModelVersionCheckpointsPager: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + + Iterating over this object will yield results and + resolve additional pages automatically. + + """ + # Create or coerce a protobuf request object. + # - Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # - Use the request object if provided (there's no risk of modifying the input as + # there are no flattened fields), or create one. + if not isinstance(request, model_service.ListModelVersionCheckpointsRequest): + request = model_service.ListModelVersionCheckpointsRequest(request) + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[ + self._transport.list_model_version_checkpoints + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # This method is paged; wrap the response in a pager, which provides + # an `__iter__` convenience method. + response = pagers.ListModelVersionCheckpointsPager( + method=rpc, + request=request, + response=response, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + def update_model( self, request: Optional[Union[model_service.UpdateModelRequest, dict]] = None, @@ -3156,16 +3316,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3211,16 +3375,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3383,16 +3551,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3504,16 +3676,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3626,16 +3802,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3686,16 +3866,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3741,16 +3925,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3796,16 +3984,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/pagers.py b/google/cloud/aiplatform_v1beta1/services/model_service/pagers.py index 330f393a33..22a880d8bd 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/pagers.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/pagers.py @@ -355,6 +355,166 @@ def __repr__(self) -> str: return "{0}<{1!r}>".format(self.__class__.__name__, self._response) +class ListModelVersionCheckpointsPager: + """A pager for iterating through ``list_model_version_checkpoints`` requests. + + This class thinly wraps an initial + :class:`google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse` object, and + provides an ``__iter__`` method to iterate through its + ``checkpoints`` field. + + If there are more pages, the ``__iter__`` method will make additional + ``ListModelVersionCheckpoints`` requests and continue to iterate + through the ``checkpoints`` field on the + corresponding responses. + + All the usual :class:`google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse` + attributes are available on the pager. If multiple requests are made, only + the most recent response is retained, and thus used for attribute lookup. + """ + + def __init__( + self, + method: Callable[..., model_service.ListModelVersionCheckpointsResponse], + request: model_service.ListModelVersionCheckpointsRequest, + response: model_service.ListModelVersionCheckpointsResponse, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = () + ): + """Instantiate the pager. + + Args: + method (Callable): The method that was originally called, and + which instantiated this pager. + request (google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest): + The initial request object. + response (google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse): + The initial response object. + retry (google.api_core.retry.Retry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + """ + self._method = method + self._request = model_service.ListModelVersionCheckpointsRequest(request) + self._response = response + self._retry = retry + self._timeout = timeout + self._metadata = metadata + + def __getattr__(self, name: str) -> Any: + return getattr(self._response, name) + + @property + def pages(self) -> Iterator[model_service.ListModelVersionCheckpointsResponse]: + yield self._response + while self._response.next_page_token: + self._request.page_token = self._response.next_page_token + self._response = self._method( + self._request, + retry=self._retry, + timeout=self._timeout, + metadata=self._metadata, + ) + yield self._response + + def __iter__(self) -> Iterator[model_service.ModelVersionCheckpoint]: + for page in self.pages: + yield from page.checkpoints + + def __repr__(self) -> str: + return "{0}<{1!r}>".format(self.__class__.__name__, self._response) + + +class ListModelVersionCheckpointsAsyncPager: + """A pager for iterating through ``list_model_version_checkpoints`` requests. + + This class thinly wraps an initial + :class:`google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse` object, and + provides an ``__aiter__`` method to iterate through its + ``checkpoints`` field. + + If there are more pages, the ``__aiter__`` method will make additional + ``ListModelVersionCheckpoints`` requests and continue to iterate + through the ``checkpoints`` field on the + corresponding responses. + + All the usual :class:`google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse` + attributes are available on the pager. If multiple requests are made, only + the most recent response is retained, and thus used for attribute lookup. + """ + + def __init__( + self, + method: Callable[ + ..., Awaitable[model_service.ListModelVersionCheckpointsResponse] + ], + request: model_service.ListModelVersionCheckpointsRequest, + response: model_service.ListModelVersionCheckpointsResponse, + *, + retry: OptionalAsyncRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = () + ): + """Instantiates the pager. + + Args: + method (Callable): The method that was originally called, and + which instantiated this pager. + request (google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest): + The initial request object. + response (google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsResponse): + The initial response object. + retry (google.api_core.retry.AsyncRetry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + """ + self._method = method + self._request = model_service.ListModelVersionCheckpointsRequest(request) + self._response = response + self._retry = retry + self._timeout = timeout + self._metadata = metadata + + def __getattr__(self, name: str) -> Any: + return getattr(self._response, name) + + @property + async def pages( + self, + ) -> AsyncIterator[model_service.ListModelVersionCheckpointsResponse]: + yield self._response + while self._response.next_page_token: + self._request.page_token = self._response.next_page_token + self._response = await self._method( + self._request, + retry=self._retry, + timeout=self._timeout, + metadata=self._metadata, + ) + yield self._response + + def __aiter__(self) -> AsyncIterator[model_service.ModelVersionCheckpoint]: + async def async_generator(): + async for page in self.pages: + for response in page.checkpoints: + yield response + + return async_generator() + + def __repr__(self) -> str: + return "{0}<{1!r}>".format(self.__class__.__name__, self._response) + + class ListModelEvaluationsPager: """A pager for iterating through ``list_model_evaluations`` requests. diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/base.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/base.py index ca2f5a50b9..c145c42c3f 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/base.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/base.py @@ -161,6 +161,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: gapic_v1.method.wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: gapic_v1.method.wrap_method( self.update_model, default_timeout=5.0, @@ -338,6 +343,18 @@ def list_model_versions( ]: raise NotImplementedError() + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + Union[ + model_service.ListModelVersionCheckpointsResponse, + Awaitable[model_service.ListModelVersionCheckpointsResponse], + ], + ]: + raise NotImplementedError() + @property def update_model( self, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc.py index f80f83e303..59220b0d76 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc.py @@ -455,6 +455,37 @@ def list_model_versions( ) return self._stubs["list_model_versions"] + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + r"""Return a callable for the list model version checkpoints method over gRPC. + + Lists checkpoints of the specified model version. + + Returns: + Callable[[~.ListModelVersionCheckpointsRequest], + ~.ListModelVersionCheckpointsResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_model_version_checkpoints" not in self._stubs: + self._stubs[ + "list_model_version_checkpoints" + ] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1beta1.ModelService/ListModelVersionCheckpoints", + request_serializer=model_service.ListModelVersionCheckpointsRequest.serialize, + response_deserializer=model_service.ListModelVersionCheckpointsResponse.deserialize, + ) + return self._stubs["list_model_version_checkpoints"] + @property def update_model( self, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc_asyncio.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc_asyncio.py index bc51b7ed75..b05e9bf371 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/grpc_asyncio.py @@ -468,6 +468,37 @@ def list_model_versions( ) return self._stubs["list_model_versions"] + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + Awaitable[model_service.ListModelVersionCheckpointsResponse], + ]: + r"""Return a callable for the list model version checkpoints method over gRPC. + + Lists checkpoints of the specified model version. + + Returns: + Callable[[~.ListModelVersionCheckpointsRequest], + Awaitable[~.ListModelVersionCheckpointsResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_model_version_checkpoints" not in self._stubs: + self._stubs[ + "list_model_version_checkpoints" + ] = self._logged_channel.unary_unary( + "/google.cloud.aiplatform.v1beta1.ModelService/ListModelVersionCheckpoints", + request_serializer=model_service.ListModelVersionCheckpointsRequest.serialize, + response_deserializer=model_service.ListModelVersionCheckpointsResponse.deserialize, + ) + return self._stubs["list_model_version_checkpoints"] + @property def update_model( self, @@ -925,6 +956,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: self._wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: self._wrap_method( self.update_model, default_timeout=5.0, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest.py index a00cabccaa..8e4ab7647c 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest.py @@ -190,6 +190,14 @@ def post_list_models(self, response): logging.log(f"Received response: {response}") return response + def pre_list_model_version_checkpoints(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_list_model_version_checkpoints(self, response): + logging.log(f"Received response: {response}") + return response + def pre_list_model_versions(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -256,12 +264,38 @@ def post_batch_import_evaluated_annotations( ) -> model_service.BatchImportEvaluatedAnnotationsResponse: """Post-rpc interceptor for batch_import_evaluated_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_evaluated_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_evaluated_annotations` interceptor runs + before the `post_batch_import_evaluated_annotations_with_metadata` interceptor. """ return response + def post_batch_import_evaluated_annotations_with_metadata( + self, + response: model_service.BatchImportEvaluatedAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportEvaluatedAnnotationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_evaluated_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_evaluated_annotations_with_metadata` + interceptor in new development instead of the `post_batch_import_evaluated_annotations` interceptor. + When both interceptors are used, this `post_batch_import_evaluated_annotations_with_metadata` interceptor runs after the + `post_batch_import_evaluated_annotations` interceptor. The (possibly modified) response returned by + `post_batch_import_evaluated_annotations` will be passed to + `post_batch_import_evaluated_annotations_with_metadata`. + """ + return response, metadata + def pre_batch_import_model_evaluation_slices( self, request: model_service.BatchImportModelEvaluationSlicesRequest, @@ -282,12 +316,38 @@ def post_batch_import_model_evaluation_slices( ) -> model_service.BatchImportModelEvaluationSlicesResponse: """Post-rpc interceptor for batch_import_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_model_evaluation_slices` interceptor runs + before the `post_batch_import_model_evaluation_slices_with_metadata` interceptor. """ return response + def post_batch_import_model_evaluation_slices_with_metadata( + self, + response: model_service.BatchImportModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_batch_import_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_batch_import_model_evaluation_slices_with_metadata` interceptor runs after the + `post_batch_import_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_batch_import_model_evaluation_slices` will be passed to + `post_batch_import_model_evaluation_slices_with_metadata`. + """ + return response, metadata + def pre_copy_model( self, request: model_service.CopyModelRequest, @@ -305,12 +365,35 @@ def post_copy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for copy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_copy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_copy_model` interceptor runs + before the `post_copy_model_with_metadata` interceptor. """ return response + def post_copy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for copy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_copy_model_with_metadata` + interceptor in new development instead of the `post_copy_model` interceptor. + When both interceptors are used, this `post_copy_model_with_metadata` interceptor runs after the + `post_copy_model` interceptor. The (possibly modified) response returned by + `post_copy_model` will be passed to + `post_copy_model_with_metadata`. + """ + return response, metadata + def pre_delete_model( self, request: model_service.DeleteModelRequest, @@ -330,12 +413,35 @@ def post_delete_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model` interceptor runs + before the `post_delete_model_with_metadata` interceptor. """ return response + def post_delete_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_with_metadata` + interceptor in new development instead of the `post_delete_model` interceptor. + When both interceptors are used, this `post_delete_model_with_metadata` interceptor runs after the + `post_delete_model` interceptor. The (possibly modified) response returned by + `post_delete_model` will be passed to + `post_delete_model_with_metadata`. + """ + return response, metadata + def pre_delete_model_version( self, request: model_service.DeleteModelVersionRequest, @@ -355,12 +461,35 @@ def post_delete_model_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_version` interceptor runs + before the `post_delete_model_version_with_metadata` interceptor. """ return response + def post_delete_model_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_version_with_metadata` + interceptor in new development instead of the `post_delete_model_version` interceptor. + When both interceptors are used, this `post_delete_model_version_with_metadata` interceptor runs after the + `post_delete_model_version` interceptor. The (possibly modified) response returned by + `post_delete_model_version` will be passed to + `post_delete_model_version_with_metadata`. + """ + return response, metadata + def pre_export_model( self, request: model_service.ExportModelRequest, @@ -380,12 +509,35 @@ def post_export_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_export_model` interceptor runs + before the `post_export_model_with_metadata` interceptor. """ return response + def post_export_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_export_model_with_metadata` + interceptor in new development instead of the `post_export_model` interceptor. + When both interceptors are used, this `post_export_model_with_metadata` interceptor runs after the + `post_export_model` interceptor. The (possibly modified) response returned by + `post_export_model` will be passed to + `post_export_model_with_metadata`. + """ + return response, metadata + def pre_get_model( self, request: model_service.GetModelRequest, @@ -401,12 +553,33 @@ def pre_get_model( def post_get_model(self, response: model.Model) -> model.Model: """Post-rpc interceptor for get_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model` interceptor runs + before the `post_get_model_with_metadata` interceptor. """ return response + def post_get_model_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_with_metadata` + interceptor in new development instead of the `post_get_model` interceptor. + When both interceptors are used, this `post_get_model_with_metadata` interceptor runs after the + `post_get_model` interceptor. The (possibly modified) response returned by + `post_get_model` will be passed to + `post_get_model_with_metadata`. + """ + return response, metadata + def pre_get_model_evaluation( self, request: model_service.GetModelEvaluationRequest, @@ -426,12 +599,37 @@ def post_get_model_evaluation( ) -> model_evaluation.ModelEvaluation: """Post-rpc interceptor for get_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation` interceptor runs + before the `post_get_model_evaluation_with_metadata` interceptor. """ return response + def post_get_model_evaluation_with_metadata( + self, + response: model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation` interceptor. + When both interceptors are used, this `post_get_model_evaluation_with_metadata` interceptor runs after the + `post_get_model_evaluation` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation` will be passed to + `post_get_model_evaluation_with_metadata`. + """ + return response, metadata + def pre_get_model_evaluation_slice( self, request: model_service.GetModelEvaluationSliceRequest, @@ -452,12 +650,38 @@ def post_get_model_evaluation_slice( ) -> model_evaluation_slice.ModelEvaluationSlice: """Post-rpc interceptor for get_model_evaluation_slice - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_slice_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation_slice` interceptor runs + before the `post_get_model_evaluation_slice_with_metadata` interceptor. """ return response + def post_get_model_evaluation_slice_with_metadata( + self, + response: model_evaluation_slice.ModelEvaluationSlice, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation_slice.ModelEvaluationSlice, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_evaluation_slice + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_slice_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation_slice` interceptor. + When both interceptors are used, this `post_get_model_evaluation_slice_with_metadata` interceptor runs after the + `post_get_model_evaluation_slice` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation_slice` will be passed to + `post_get_model_evaluation_slice_with_metadata`. + """ + return response, metadata + def pre_import_model_evaluation( self, request: model_service.ImportModelEvaluationRequest, @@ -478,12 +702,37 @@ def post_import_model_evaluation( ) -> gca_model_evaluation.ModelEvaluation: """Post-rpc interceptor for import_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_import_model_evaluation` interceptor runs + before the `post_import_model_evaluation_with_metadata` interceptor. """ return response + def post_import_model_evaluation_with_metadata( + self, + response: gca_model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for import_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_import_model_evaluation_with_metadata` + interceptor in new development instead of the `post_import_model_evaluation` interceptor. + When both interceptors are used, this `post_import_model_evaluation_with_metadata` interceptor runs after the + `post_import_model_evaluation` interceptor. The (possibly modified) response returned by + `post_import_model_evaluation` will be passed to + `post_import_model_evaluation_with_metadata`. + """ + return response, metadata + def pre_list_model_evaluations( self, request: model_service.ListModelEvaluationsRequest, @@ -504,12 +753,38 @@ def post_list_model_evaluations( ) -> model_service.ListModelEvaluationsResponse: """Post-rpc interceptor for list_model_evaluations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluations` interceptor runs + before the `post_list_model_evaluations_with_metadata` interceptor. """ return response + def post_list_model_evaluations_with_metadata( + self, + response: model_service.ListModelEvaluationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluations_with_metadata` + interceptor in new development instead of the `post_list_model_evaluations` interceptor. + When both interceptors are used, this `post_list_model_evaluations_with_metadata` interceptor runs after the + `post_list_model_evaluations` interceptor. The (possibly modified) response returned by + `post_list_model_evaluations` will be passed to + `post_list_model_evaluations_with_metadata`. + """ + return response, metadata + def pre_list_model_evaluation_slices( self, request: model_service.ListModelEvaluationSlicesRequest, @@ -530,12 +805,38 @@ def post_list_model_evaluation_slices( ) -> model_service.ListModelEvaluationSlicesResponse: """Post-rpc interceptor for list_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluation_slices` interceptor runs + before the `post_list_model_evaluation_slices_with_metadata` interceptor. """ return response + def post_list_model_evaluation_slices_with_metadata( + self, + response: model_service.ListModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_list_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_list_model_evaluation_slices_with_metadata` interceptor runs after the + `post_list_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_list_model_evaluation_slices` will be passed to + `post_list_model_evaluation_slices_with_metadata`. + """ + return response, metadata + def pre_list_models( self, request: model_service.ListModelsRequest, @@ -555,12 +856,89 @@ def post_list_models( ) -> model_service.ListModelsResponse: """Post-rpc interceptor for list_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_models` interceptor runs + before the `post_list_models_with_metadata` interceptor. """ return response + def post_list_models_with_metadata( + self, + response: model_service.ListModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_models_with_metadata` + interceptor in new development instead of the `post_list_models` interceptor. + When both interceptors are used, this `post_list_models_with_metadata` interceptor runs after the + `post_list_models` interceptor. The (possibly modified) response returned by + `post_list_models` will be passed to + `post_list_models_with_metadata`. + """ + return response, metadata + + def pre_list_model_version_checkpoints( + self, + request: model_service.ListModelVersionCheckpointsRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to manipulate the request or metadata + before they are sent to the ModelService server. + """ + return request, metadata + + def post_list_model_version_checkpoints( + self, response: model_service.ListModelVersionCheckpointsResponse + ) -> model_service.ListModelVersionCheckpointsResponse: + """Post-rpc interceptor for list_model_version_checkpoints + + DEPRECATED. Please use the `post_list_model_version_checkpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the ModelService server but before + it is returned to user code. This `post_list_model_version_checkpoints` interceptor runs + before the `post_list_model_version_checkpoints_with_metadata` interceptor. + """ + return response + + def post_list_model_version_checkpoints_with_metadata( + self, + response: model_service.ListModelVersionCheckpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_version_checkpoints_with_metadata` + interceptor in new development instead of the `post_list_model_version_checkpoints` interceptor. + When both interceptors are used, this `post_list_model_version_checkpoints_with_metadata` interceptor runs after the + `post_list_model_version_checkpoints` interceptor. The (possibly modified) response returned by + `post_list_model_version_checkpoints` will be passed to + `post_list_model_version_checkpoints_with_metadata`. + """ + return response, metadata + def pre_list_model_versions( self, request: model_service.ListModelVersionsRequest, @@ -580,12 +958,37 @@ def post_list_model_versions( ) -> model_service.ListModelVersionsResponse: """Post-rpc interceptor for list_model_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_versions` interceptor runs + before the `post_list_model_versions_with_metadata` interceptor. """ return response + def post_list_model_versions_with_metadata( + self, + response: model_service.ListModelVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_model_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_versions_with_metadata` + interceptor in new development instead of the `post_list_model_versions` interceptor. + When both interceptors are used, this `post_list_model_versions_with_metadata` interceptor runs after the + `post_list_model_versions` interceptor. The (possibly modified) response returned by + `post_list_model_versions` will be passed to + `post_list_model_versions_with_metadata`. + """ + return response, metadata + def pre_merge_version_aliases( self, request: model_service.MergeVersionAliasesRequest, @@ -604,12 +1007,33 @@ def pre_merge_version_aliases( def post_merge_version_aliases(self, response: model.Model) -> model.Model: """Post-rpc interceptor for merge_version_aliases - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_merge_version_aliases_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_merge_version_aliases` interceptor runs + before the `post_merge_version_aliases_with_metadata` interceptor. """ return response + def post_merge_version_aliases_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for merge_version_aliases + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_merge_version_aliases_with_metadata` + interceptor in new development instead of the `post_merge_version_aliases` interceptor. + When both interceptors are used, this `post_merge_version_aliases_with_metadata` interceptor runs after the + `post_merge_version_aliases` interceptor. The (possibly modified) response returned by + `post_merge_version_aliases` will be passed to + `post_merge_version_aliases_with_metadata`. + """ + return response, metadata + def pre_update_explanation_dataset( self, request: model_service.UpdateExplanationDatasetRequest, @@ -630,12 +1054,35 @@ def post_update_explanation_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_explanation_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_explanation_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_explanation_dataset` interceptor runs + before the `post_update_explanation_dataset_with_metadata` interceptor. """ return response + def post_update_explanation_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_explanation_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_explanation_dataset_with_metadata` + interceptor in new development instead of the `post_update_explanation_dataset` interceptor. + When both interceptors are used, this `post_update_explanation_dataset_with_metadata` interceptor runs after the + `post_update_explanation_dataset` interceptor. The (possibly modified) response returned by + `post_update_explanation_dataset` will be passed to + `post_update_explanation_dataset_with_metadata`. + """ + return response, metadata + def pre_update_model( self, request: model_service.UpdateModelRequest, @@ -653,12 +1100,35 @@ def pre_update_model( def post_update_model(self, response: gca_model.Model) -> gca_model.Model: """Post-rpc interceptor for update_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_model` interceptor runs + before the `post_update_model_with_metadata` interceptor. """ return response + def post_update_model_with_metadata( + self, + response: gca_model.Model, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_model_with_metadata` + interceptor in new development instead of the `post_update_model` interceptor. + When both interceptors are used, this `post_update_model_with_metadata` interceptor runs after the + `post_update_model` interceptor. The (possibly modified) response returned by + `post_update_model` will be passed to + `post_update_model_with_metadata`. + """ + return response, metadata + def pre_upload_model( self, request: model_service.UploadModelRequest, @@ -678,12 +1148,35 @@ def post_upload_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for upload_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_upload_model` interceptor runs + before the `post_upload_model_with_metadata` interceptor. """ return response + def post_upload_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upload_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_upload_model_with_metadata` + interceptor in new development instead of the `post_upload_model` interceptor. + When both interceptors are used, this `post_upload_model_with_metadata` interceptor runs after the + `post_upload_model` interceptor. The (possibly modified) response returned by + `post_upload_model` will be passed to + `post_upload_model_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3235,6 +3728,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_import_evaluated_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_import_evaluated_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3396,6 +3896,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_import_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_import_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3553,6 +4060,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_copy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_copy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3700,6 +4211,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3845,6 +4360,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_model_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_model_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4000,6 +4519,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_export_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4144,6 +4667,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4292,6 +4819,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4444,6 +4975,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_model_evaluation_slice(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_model_evaluation_slice_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4600,6 +5135,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4748,6 +5287,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_evaluations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_evaluations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4900,6 +5443,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5052,6 +5599,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5077,6 +5628,167 @@ def __call__( ) return resp + class _ListModelVersionCheckpoints( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints, + ModelServiceRestStub, + ): + def __hash__(self): + return hash("ModelServiceRestTransport.ListModelVersionCheckpoints") + + @staticmethod + def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + ) + return response + + def __call__( + self, + request: model_service.ListModelVersionCheckpointsRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> model_service.ListModelVersionCheckpointsResponse: + r"""Call the list model version + checkpoints method over HTTP. + + Args: + request (~.model_service.ListModelVersionCheckpointsRequest): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.model_service.ListModelVersionCheckpointsResponse: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + + """ + + http_options = ( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_http_options() + ) + + request, metadata = self._interceptor.pre_list_model_version_checkpoints( + request, metadata + ) + transcoded_request = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_transcoded_request( + http_options, request + ) + + # Jsonify the query params + query_params = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = type(request).to_json(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1beta1.ModelServiceClient.ListModelVersionCheckpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = ( + ModelServiceRestTransport._ListModelVersionCheckpoints._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + ) + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = model_service.ListModelVersionCheckpointsResponse() + pb_resp = model_service.ListModelVersionCheckpointsResponse.pb(resp) + + json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) + + resp = self._interceptor.post_list_model_version_checkpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_model_version_checkpoints_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = ( + model_service.ListModelVersionCheckpointsResponse.to_json( + response + ) + ) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": response.status_code, + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1beta1.ModelServiceClient.list_model_version_checkpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + return resp + class _ListModelVersions( _BaseModelServiceRestTransport._BaseListModelVersions, ModelServiceRestStub ): @@ -5200,6 +5912,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_model_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_model_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5352,6 +6068,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_merge_version_aliases(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_merge_version_aliases_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5507,6 +6227,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_explanation_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_explanation_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5661,6 +6385,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5816,6 +6544,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_upload_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upload_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5961,6 +6693,17 @@ def list_models( # In C++ this would require a dynamic_cast return self._ListModels(self._session, self._host, self._interceptor) # type: ignore + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._ListModelVersionCheckpoints(self._session, self._host, self._interceptor) # type: ignore + @property def list_model_versions( self, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_asyncio.py index 5bbb83b666..07061f2e9c 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_asyncio.py @@ -208,6 +208,14 @@ async def post_list_models(self, response): logging.log(f"Received response: {response}") return response + async def pre_list_model_version_checkpoints(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + async def post_list_model_version_checkpoints(self, response): + logging.log(f"Received response: {response}") + return response + async def pre_list_model_versions(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -274,12 +282,38 @@ async def post_batch_import_evaluated_annotations( ) -> model_service.BatchImportEvaluatedAnnotationsResponse: """Post-rpc interceptor for batch_import_evaluated_annotations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_evaluated_annotations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_evaluated_annotations` interceptor runs + before the `post_batch_import_evaluated_annotations_with_metadata` interceptor. """ return response + async def post_batch_import_evaluated_annotations_with_metadata( + self, + response: model_service.BatchImportEvaluatedAnnotationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportEvaluatedAnnotationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_evaluated_annotations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_evaluated_annotations_with_metadata` + interceptor in new development instead of the `post_batch_import_evaluated_annotations` interceptor. + When both interceptors are used, this `post_batch_import_evaluated_annotations_with_metadata` interceptor runs after the + `post_batch_import_evaluated_annotations` interceptor. The (possibly modified) response returned by + `post_batch_import_evaluated_annotations` will be passed to + `post_batch_import_evaluated_annotations_with_metadata`. + """ + return response, metadata + async def pre_batch_import_model_evaluation_slices( self, request: model_service.BatchImportModelEvaluationSlicesRequest, @@ -300,12 +334,38 @@ async def post_batch_import_model_evaluation_slices( ) -> model_service.BatchImportModelEvaluationSlicesResponse: """Post-rpc interceptor for batch_import_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_import_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_batch_import_model_evaluation_slices` interceptor runs + before the `post_batch_import_model_evaluation_slices_with_metadata` interceptor. """ return response + async def post_batch_import_model_evaluation_slices_with_metadata( + self, + response: model_service.BatchImportModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.BatchImportModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_import_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_batch_import_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_batch_import_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_batch_import_model_evaluation_slices_with_metadata` interceptor runs after the + `post_batch_import_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_batch_import_model_evaluation_slices` will be passed to + `post_batch_import_model_evaluation_slices_with_metadata`. + """ + return response, metadata + async def pre_copy_model( self, request: model_service.CopyModelRequest, @@ -323,12 +383,35 @@ async def post_copy_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for copy_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_copy_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_copy_model` interceptor runs + before the `post_copy_model_with_metadata` interceptor. """ return response + async def post_copy_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for copy_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_copy_model_with_metadata` + interceptor in new development instead of the `post_copy_model` interceptor. + When both interceptors are used, this `post_copy_model_with_metadata` interceptor runs after the + `post_copy_model` interceptor. The (possibly modified) response returned by + `post_copy_model` will be passed to + `post_copy_model_with_metadata`. + """ + return response, metadata + async def pre_delete_model( self, request: model_service.DeleteModelRequest, @@ -348,12 +431,35 @@ async def post_delete_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model` interceptor runs + before the `post_delete_model_with_metadata` interceptor. """ return response + async def post_delete_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_with_metadata` + interceptor in new development instead of the `post_delete_model` interceptor. + When both interceptors are used, this `post_delete_model_with_metadata` interceptor runs after the + `post_delete_model` interceptor. The (possibly modified) response returned by + `post_delete_model` will be passed to + `post_delete_model_with_metadata`. + """ + return response, metadata + async def pre_delete_model_version( self, request: model_service.DeleteModelVersionRequest, @@ -373,12 +479,35 @@ async def post_delete_model_version( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_model_version - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_model_version_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_delete_model_version` interceptor runs + before the `post_delete_model_version_with_metadata` interceptor. """ return response + async def post_delete_model_version_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_model_version + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_delete_model_version_with_metadata` + interceptor in new development instead of the `post_delete_model_version` interceptor. + When both interceptors are used, this `post_delete_model_version_with_metadata` interceptor runs after the + `post_delete_model_version` interceptor. The (possibly modified) response returned by + `post_delete_model_version` will be passed to + `post_delete_model_version_with_metadata`. + """ + return response, metadata + async def pre_export_model( self, request: model_service.ExportModelRequest, @@ -398,12 +527,35 @@ async def post_export_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for export_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_export_model` interceptor runs + before the `post_export_model_with_metadata` interceptor. """ return response + async def post_export_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for export_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_export_model_with_metadata` + interceptor in new development instead of the `post_export_model` interceptor. + When both interceptors are used, this `post_export_model_with_metadata` interceptor runs after the + `post_export_model` interceptor. The (possibly modified) response returned by + `post_export_model` will be passed to + `post_export_model_with_metadata`. + """ + return response, metadata + async def pre_get_model( self, request: model_service.GetModelRequest, @@ -419,12 +571,33 @@ async def pre_get_model( async def post_get_model(self, response: model.Model) -> model.Model: """Post-rpc interceptor for get_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model` interceptor runs + before the `post_get_model_with_metadata` interceptor. """ return response + async def post_get_model_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_with_metadata` + interceptor in new development instead of the `post_get_model` interceptor. + When both interceptors are used, this `post_get_model_with_metadata` interceptor runs after the + `post_get_model` interceptor. The (possibly modified) response returned by + `post_get_model` will be passed to + `post_get_model_with_metadata`. + """ + return response, metadata + async def pre_get_model_evaluation( self, request: model_service.GetModelEvaluationRequest, @@ -444,12 +617,37 @@ async def post_get_model_evaluation( ) -> model_evaluation.ModelEvaluation: """Post-rpc interceptor for get_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation` interceptor runs + before the `post_get_model_evaluation_with_metadata` interceptor. """ return response + async def post_get_model_evaluation_with_metadata( + self, + response: model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation` interceptor. + When both interceptors are used, this `post_get_model_evaluation_with_metadata` interceptor runs after the + `post_get_model_evaluation` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation` will be passed to + `post_get_model_evaluation_with_metadata`. + """ + return response, metadata + async def pre_get_model_evaluation_slice( self, request: model_service.GetModelEvaluationSliceRequest, @@ -470,12 +668,38 @@ async def post_get_model_evaluation_slice( ) -> model_evaluation_slice.ModelEvaluationSlice: """Post-rpc interceptor for get_model_evaluation_slice - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_model_evaluation_slice_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_get_model_evaluation_slice` interceptor runs + before the `post_get_model_evaluation_slice_with_metadata` interceptor. """ return response + async def post_get_model_evaluation_slice_with_metadata( + self, + response: model_evaluation_slice.ModelEvaluationSlice, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_evaluation_slice.ModelEvaluationSlice, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_model_evaluation_slice + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_get_model_evaluation_slice_with_metadata` + interceptor in new development instead of the `post_get_model_evaluation_slice` interceptor. + When both interceptors are used, this `post_get_model_evaluation_slice_with_metadata` interceptor runs after the + `post_get_model_evaluation_slice` interceptor. The (possibly modified) response returned by + `post_get_model_evaluation_slice` will be passed to + `post_get_model_evaluation_slice_with_metadata`. + """ + return response, metadata + async def pre_import_model_evaluation( self, request: model_service.ImportModelEvaluationRequest, @@ -496,12 +720,37 @@ async def post_import_model_evaluation( ) -> gca_model_evaluation.ModelEvaluation: """Post-rpc interceptor for import_model_evaluation - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_model_evaluation_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_import_model_evaluation` interceptor runs + before the `post_import_model_evaluation_with_metadata` interceptor. """ return response + async def post_import_model_evaluation_with_metadata( + self, + response: gca_model_evaluation.ModelEvaluation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_model_evaluation.ModelEvaluation, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for import_model_evaluation + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_import_model_evaluation_with_metadata` + interceptor in new development instead of the `post_import_model_evaluation` interceptor. + When both interceptors are used, this `post_import_model_evaluation_with_metadata` interceptor runs after the + `post_import_model_evaluation` interceptor. The (possibly modified) response returned by + `post_import_model_evaluation` will be passed to + `post_import_model_evaluation_with_metadata`. + """ + return response, metadata + async def pre_list_model_evaluations( self, request: model_service.ListModelEvaluationsRequest, @@ -522,12 +771,38 @@ async def post_list_model_evaluations( ) -> model_service.ListModelEvaluationsResponse: """Post-rpc interceptor for list_model_evaluations - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluations_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluations` interceptor runs + before the `post_list_model_evaluations_with_metadata` interceptor. """ return response + async def post_list_model_evaluations_with_metadata( + self, + response: model_service.ListModelEvaluationsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluations + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluations_with_metadata` + interceptor in new development instead of the `post_list_model_evaluations` interceptor. + When both interceptors are used, this `post_list_model_evaluations_with_metadata` interceptor runs after the + `post_list_model_evaluations` interceptor. The (possibly modified) response returned by + `post_list_model_evaluations` will be passed to + `post_list_model_evaluations_with_metadata`. + """ + return response, metadata + async def pre_list_model_evaluation_slices( self, request: model_service.ListModelEvaluationSlicesRequest, @@ -548,12 +823,38 @@ async def post_list_model_evaluation_slices( ) -> model_service.ListModelEvaluationSlicesResponse: """Post-rpc interceptor for list_model_evaluation_slices - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_evaluation_slices_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_evaluation_slices` interceptor runs + before the `post_list_model_evaluation_slices_with_metadata` interceptor. """ return response + async def post_list_model_evaluation_slices_with_metadata( + self, + response: model_service.ListModelEvaluationSlicesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelEvaluationSlicesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_evaluation_slices + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_evaluation_slices_with_metadata` + interceptor in new development instead of the `post_list_model_evaluation_slices` interceptor. + When both interceptors are used, this `post_list_model_evaluation_slices_with_metadata` interceptor runs after the + `post_list_model_evaluation_slices` interceptor. The (possibly modified) response returned by + `post_list_model_evaluation_slices` will be passed to + `post_list_model_evaluation_slices_with_metadata`. + """ + return response, metadata + async def pre_list_models( self, request: model_service.ListModelsRequest, @@ -573,12 +874,89 @@ async def post_list_models( ) -> model_service.ListModelsResponse: """Post-rpc interceptor for list_models - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_models_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_models` interceptor runs + before the `post_list_models_with_metadata` interceptor. + """ + return response + + async def post_list_models_with_metadata( + self, + response: model_service.ListModelsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_models + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_models_with_metadata` + interceptor in new development instead of the `post_list_models` interceptor. + When both interceptors are used, this `post_list_models_with_metadata` interceptor runs after the + `post_list_models` interceptor. The (possibly modified) response returned by + `post_list_models` will be passed to + `post_list_models_with_metadata`. + """ + return response, metadata + + async def pre_list_model_version_checkpoints( + self, + request: model_service.ListModelVersionCheckpointsRequest, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsRequest, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Pre-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to manipulate the request or metadata + before they are sent to the ModelService server. + """ + return request, metadata + + async def post_list_model_version_checkpoints( + self, response: model_service.ListModelVersionCheckpointsResponse + ) -> model_service.ListModelVersionCheckpointsResponse: + """Post-rpc interceptor for list_model_version_checkpoints + + DEPRECATED. Please use the `post_list_model_version_checkpoints_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response + after it is returned by the ModelService server but before + it is returned to user code. This `post_list_model_version_checkpoints` interceptor runs + before the `post_list_model_version_checkpoints_with_metadata` interceptor. """ return response + async def post_list_model_version_checkpoints_with_metadata( + self, + response: model_service.ListModelVersionCheckpointsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionCheckpointsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_model_version_checkpoints + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_version_checkpoints_with_metadata` + interceptor in new development instead of the `post_list_model_version_checkpoints` interceptor. + When both interceptors are used, this `post_list_model_version_checkpoints_with_metadata` interceptor runs after the + `post_list_model_version_checkpoints` interceptor. The (possibly modified) response returned by + `post_list_model_version_checkpoints` will be passed to + `post_list_model_version_checkpoints_with_metadata`. + """ + return response, metadata + async def pre_list_model_versions( self, request: model_service.ListModelVersionsRequest, @@ -598,12 +976,37 @@ async def post_list_model_versions( ) -> model_service.ListModelVersionsResponse: """Post-rpc interceptor for list_model_versions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_model_versions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_list_model_versions` interceptor runs + before the `post_list_model_versions_with_metadata` interceptor. """ return response + async def post_list_model_versions_with_metadata( + self, + response: model_service.ListModelVersionsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + model_service.ListModelVersionsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_model_versions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_list_model_versions_with_metadata` + interceptor in new development instead of the `post_list_model_versions` interceptor. + When both interceptors are used, this `post_list_model_versions_with_metadata` interceptor runs after the + `post_list_model_versions` interceptor. The (possibly modified) response returned by + `post_list_model_versions` will be passed to + `post_list_model_versions_with_metadata`. + """ + return response, metadata + async def pre_merge_version_aliases( self, request: model_service.MergeVersionAliasesRequest, @@ -622,12 +1025,33 @@ async def pre_merge_version_aliases( async def post_merge_version_aliases(self, response: model.Model) -> model.Model: """Post-rpc interceptor for merge_version_aliases - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_merge_version_aliases_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_merge_version_aliases` interceptor runs + before the `post_merge_version_aliases_with_metadata` interceptor. """ return response + async def post_merge_version_aliases_with_metadata( + self, response: model.Model, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for merge_version_aliases + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_merge_version_aliases_with_metadata` + interceptor in new development instead of the `post_merge_version_aliases` interceptor. + When both interceptors are used, this `post_merge_version_aliases_with_metadata` interceptor runs after the + `post_merge_version_aliases` interceptor. The (possibly modified) response returned by + `post_merge_version_aliases` will be passed to + `post_merge_version_aliases_with_metadata`. + """ + return response, metadata + async def pre_update_explanation_dataset( self, request: model_service.UpdateExplanationDatasetRequest, @@ -648,12 +1072,35 @@ async def post_update_explanation_dataset( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_explanation_dataset - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_explanation_dataset_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_explanation_dataset` interceptor runs + before the `post_update_explanation_dataset_with_metadata` interceptor. """ return response + async def post_update_explanation_dataset_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_explanation_dataset + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_explanation_dataset_with_metadata` + interceptor in new development instead of the `post_update_explanation_dataset` interceptor. + When both interceptors are used, this `post_update_explanation_dataset_with_metadata` interceptor runs after the + `post_update_explanation_dataset` interceptor. The (possibly modified) response returned by + `post_update_explanation_dataset` will be passed to + `post_update_explanation_dataset_with_metadata`. + """ + return response, metadata + async def pre_update_model( self, request: model_service.UpdateModelRequest, @@ -671,12 +1118,35 @@ async def pre_update_model( async def post_update_model(self, response: gca_model.Model) -> gca_model.Model: """Post-rpc interceptor for update_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_update_model` interceptor runs + before the `post_update_model_with_metadata` interceptor. """ return response + async def post_update_model_with_metadata( + self, + response: gca_model.Model, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_model.Model, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_update_model_with_metadata` + interceptor in new development instead of the `post_update_model` interceptor. + When both interceptors are used, this `post_update_model_with_metadata` interceptor runs after the + `post_update_model` interceptor. The (possibly modified) response returned by + `post_update_model` will be passed to + `post_update_model_with_metadata`. + """ + return response, metadata + async def pre_upload_model( self, request: model_service.UploadModelRequest, @@ -696,12 +1166,35 @@ async def post_upload_model( ) -> operations_pb2.Operation: """Post-rpc interceptor for upload_model - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_model_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ModelService server but before - it is returned to user code. + it is returned to user code. This `post_upload_model` interceptor runs + before the `post_upload_model_with_metadata` interceptor. """ return response + async def post_upload_model_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upload_model + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ModelService server but before it is returned to user code. + + We recommend only using this `post_upload_model_with_metadata` + interceptor in new development instead of the `post_upload_model` interceptor. + When both interceptors are used, this `post_upload_model_with_metadata` interceptor runs after the + `post_upload_model` interceptor. The (possibly modified) response returned by + `post_upload_model` will be passed to + `post_upload_model_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1037,6 +1530,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.list_model_version_checkpoints: self._wrap_method( + self.list_model_version_checkpoints, + default_timeout=None, + client_info=client_info, + ), self.update_model: self._wrap_method( self.update_model, default_timeout=5.0, @@ -1305,6 +1803,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_import_evaluated_annotations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_import_evaluated_annotations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1476,6 +1981,13 @@ async def __call__( resp = await self._interceptor.post_batch_import_model_evaluation_slices( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_import_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1643,6 +2155,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_copy_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_copy_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1800,6 +2316,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1956,6 +2476,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_model_version(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_model_version_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2121,6 +2645,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_export_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_export_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2273,6 +2801,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2430,6 +2962,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2588,6 +3124,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_model_evaluation_slice(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_model_evaluation_slice_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2751,6 +3294,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_model_evaluation(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_import_model_evaluation_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2906,6 +3456,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_evaluations(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_model_evaluations_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3065,6 +3619,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_evaluation_slices(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_evaluation_slices_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3225,6 +3786,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_models(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_models_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3251,6 +3816,174 @@ async def __call__( return resp + class _ListModelVersionCheckpoints( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints, + AsyncModelServiceRestStub, + ): + def __hash__(self): + return hash("AsyncModelServiceRestTransport.ListModelVersionCheckpoints") + + @staticmethod + async def _get_response( + host, + metadata, + query_params, + session, + timeout, + transcoded_request, + body=None, + ): + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = await getattr(session, method)( + "{host}{uri}".format(host=host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + ) + return response + + async def __call__( + self, + request: model_service.ListModelVersionCheckpointsRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, Union[str, bytes]]] = (), + ) -> model_service.ListModelVersionCheckpointsResponse: + r"""Call the list model version + checkpoints method over HTTP. + + Args: + request (~.model_service.ListModelVersionCheckpointsRequest): + The request object. Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be + sent along with the request as metadata. Normally, each value must be of type `str`, + but for metadata keys ending with the suffix `-bin`, the corresponding values must + be of type `bytes`. + + Returns: + ~.model_service.ListModelVersionCheckpointsResponse: + Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + + """ + + http_options = ( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_http_options() + ) + + ( + request, + metadata, + ) = await self._interceptor.pre_list_model_version_checkpoints( + request, metadata + ) + transcoded_request = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_transcoded_request( + http_options, request + ) + + # Jsonify the query params + query_params = _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_query_params_json( + transcoded_request + ) + + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + try: + request_payload = type(request).to_json(request) + except: + request_payload = None + http_request = { + "payload": request_payload, + "requestMethod": method, + "requestUrl": request_url, + "headers": dict(metadata), + } + _LOGGER.debug( + f"Sending request for google.cloud.aiplatform_v1beta1.ModelServiceClient.ListModelVersionCheckpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "httpRequest": http_request, + "metadata": http_request["headers"], + }, + ) + + # Send the request + response = await AsyncModelServiceRestTransport._ListModelVersionCheckpoints._get_response( + self._host, + metadata, + query_params, + self._session, + timeout, + transcoded_request, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + content = await response.read() + payload = json.loads(content.decode("utf-8")) + request_url = "{host}{uri}".format( + host=self._host, uri=transcoded_request["uri"] + ) + method = transcoded_request["method"] + raise core_exceptions.format_http_response_error(response, method, request_url, payload) # type: ignore + + # Return the response + resp = model_service.ListModelVersionCheckpointsResponse() + pb_resp = model_service.ListModelVersionCheckpointsResponse.pb(resp) + content = await response.read() + json_format.Parse(content, pb_resp, ignore_unknown_fields=True) + resp = await self._interceptor.post_list_model_version_checkpoints(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_model_version_checkpoints_with_metadata( + resp, response_metadata + ) + if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( + logging.DEBUG + ): # pragma: NO COVER + try: + response_payload = ( + model_service.ListModelVersionCheckpointsResponse.to_json( + response + ) + ) + except: + response_payload = None + http_response = { + "payload": response_payload, + "headers": dict(response.headers), + "status": "OK", # need to obtain this properly + } + _LOGGER.debug( + "Received response for google.cloud.aiplatform_v1beta1.ModelServiceAsyncClient.list_model_version_checkpoints", + extra={ + "serviceName": "google.cloud.aiplatform.v1beta1.ModelService", + "rpcName": "ListModelVersionCheckpoints", + "metadata": http_response["headers"], + "httpResponse": http_response, + }, + ) + + return resp + class _ListModelVersions( _BaseModelServiceRestTransport._BaseListModelVersions, AsyncModelServiceRestStub ): @@ -3381,6 +4114,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_model_versions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_model_versions_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3542,6 +4279,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_merge_version_aliases(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_merge_version_aliases_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3703,6 +4444,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_explanation_dataset(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_explanation_dataset_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3865,6 +4613,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4030,6 +4782,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upload_model(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upload_model_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6235,6 +6991,15 @@ def list_models( ) -> Callable[[model_service.ListModelsRequest], model_service.ListModelsResponse]: return self._ListModels(self._session, self._host, self._interceptor) # type: ignore + @property + def list_model_version_checkpoints( + self, + ) -> Callable[ + [model_service.ListModelVersionCheckpointsRequest], + model_service.ListModelVersionCheckpointsResponse, + ]: + return self._ListModelVersionCheckpoints(self._session, self._host, self._interceptor) # type: ignore + @property def list_model_versions( self, diff --git a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_base.py b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_base.py index ecb084f7ef..ed0a347c76 100644 --- a/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_base.py +++ b/google/cloud/aiplatform_v1beta1/services/model_service/transports/rest_base.py @@ -765,6 +765,53 @@ def _get_query_params_json(transcoded_request): query_params["$alt"] = "json;enum-encoding=int" return query_params + class _BaseListModelVersionCheckpoints: + def __hash__(self): # pragma: NO COVER + return NotImplementedError("__hash__ must be implemented.") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {} + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + @staticmethod + def _get_http_options(): + http_options: List[Dict[str, str]] = [ + { + "method": "get", + "uri": "/v1beta1/{name=projects/*/locations/*/models/*}:listCheckpoints", + }, + ] + return http_options + + @staticmethod + def _get_transcoded_request(http_options, request): + pb_request = model_service.ListModelVersionCheckpointsRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + return transcoded_request + + @staticmethod + def _get_query_params_json(transcoded_request): + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update( + _BaseModelServiceRestTransport._BaseListModelVersionCheckpoints._get_unset_required_fields( + query_params + ) + ) + + query_params["$alt"] = "json;enum-encoding=int" + return query_params + class _BaseListModelVersions: def __hash__(self): # pragma: NO COVER return NotImplementedError("__hash__ must be implemented.") diff --git a/google/cloud/aiplatform_v1beta1/services/notebook_service/async_client.py b/google/cloud/aiplatform_v1beta1/services/notebook_service/async_client.py index 486f62d471..44f6d34f27 100644 --- a/google/cloud/aiplatform_v1beta1/services/notebook_service/async_client.py +++ b/google/cloud/aiplatform_v1beta1/services/notebook_service/async_client.py @@ -63,6 +63,7 @@ ) from google.cloud.aiplatform_v1beta1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1beta1.types import notebook_service +from google.cloud.aiplatform_v1beta1.types import notebook_software_config from google.cloud.aiplatform_v1beta1.types import operation as gca_operation from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore diff --git a/google/cloud/aiplatform_v1beta1/services/notebook_service/client.py b/google/cloud/aiplatform_v1beta1/services/notebook_service/client.py index 8b42939a21..a3bb2e4e96 100644 --- a/google/cloud/aiplatform_v1beta1/services/notebook_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/notebook_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -77,6 +79,7 @@ ) from google.cloud.aiplatform_v1beta1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1beta1.types import notebook_service +from google.cloud.aiplatform_v1beta1.types import notebook_software_config from google.cloud.aiplatform_v1beta1.types import operation as gca_operation from google.cloud.location import locations_pb2 # type: ignore from google.iam.v1 import iam_policy_pb2 # type: ignore @@ -655,6 +658,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -3044,16 +3074,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -3099,16 +3133,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -3271,16 +3309,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3392,16 +3434,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3514,16 +3560,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3574,16 +3624,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3629,16 +3683,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3684,16 +3742,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest.py index 94c752d0f3..680071cecb 100644 --- a/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest.py @@ -235,12 +235,35 @@ def post_assign_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for assign_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_assign_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_assign_notebook_runtime` interceptor runs + before the `post_assign_notebook_runtime_with_metadata` interceptor. """ return response + def post_assign_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for assign_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_assign_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_assign_notebook_runtime` interceptor. + When both interceptors are used, this `post_assign_notebook_runtime_with_metadata` interceptor runs after the + `post_assign_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_assign_notebook_runtime` will be passed to + `post_assign_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_create_notebook_execution_job( self, request: notebook_service.CreateNotebookExecutionJobRequest, @@ -261,12 +284,35 @@ def post_create_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_execution_job` interceptor runs + before the `post_create_notebook_execution_job_with_metadata` interceptor. """ return response + def post_create_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_create_notebook_execution_job` interceptor. + When both interceptors are used, this `post_create_notebook_execution_job_with_metadata` interceptor runs after the + `post_create_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_create_notebook_execution_job` will be passed to + `post_create_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_create_notebook_runtime_template( self, request: notebook_service.CreateNotebookRuntimeTemplateRequest, @@ -287,12 +333,35 @@ def post_create_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_runtime_template` interceptor runs + before the `post_create_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_create_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_create_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_create_notebook_runtime_template_with_metadata` interceptor runs after the + `post_create_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_create_notebook_runtime_template` will be passed to + `post_create_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_execution_job( self, request: notebook_service.DeleteNotebookExecutionJobRequest, @@ -313,12 +382,35 @@ def post_delete_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_execution_job` interceptor runs + before the `post_delete_notebook_execution_job_with_metadata` interceptor. """ return response + def post_delete_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_delete_notebook_execution_job` interceptor. + When both interceptors are used, this `post_delete_notebook_execution_job_with_metadata` interceptor runs after the + `post_delete_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_delete_notebook_execution_job` will be passed to + `post_delete_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_runtime( self, request: notebook_service.DeleteNotebookRuntimeRequest, @@ -339,12 +431,35 @@ def post_delete_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime` interceptor runs + before the `post_delete_notebook_runtime_with_metadata` interceptor. """ return response + def post_delete_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_with_metadata` interceptor runs after the + `post_delete_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime` will be passed to + `post_delete_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_delete_notebook_runtime_template( self, request: notebook_service.DeleteNotebookRuntimeTemplateRequest, @@ -365,12 +480,35 @@ def post_delete_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime_template` interceptor runs + before the `post_delete_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_delete_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_template_with_metadata` interceptor runs after the + `post_delete_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime_template` will be passed to + `post_delete_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_get_notebook_execution_job( self, request: notebook_service.GetNotebookExecutionJobRequest, @@ -391,12 +529,38 @@ def post_get_notebook_execution_job( ) -> notebook_execution_job.NotebookExecutionJob: """Post-rpc interceptor for get_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_execution_job` interceptor runs + before the `post_get_notebook_execution_job_with_metadata` interceptor. """ return response + def post_get_notebook_execution_job_with_metadata( + self, + response: notebook_execution_job.NotebookExecutionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_execution_job.NotebookExecutionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_get_notebook_execution_job` interceptor. + When both interceptors are used, this `post_get_notebook_execution_job_with_metadata` interceptor runs after the + `post_get_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_get_notebook_execution_job` will be passed to + `post_get_notebook_execution_job_with_metadata`. + """ + return response, metadata + def pre_get_notebook_runtime( self, request: notebook_service.GetNotebookRuntimeRequest, @@ -417,12 +581,37 @@ def post_get_notebook_runtime( ) -> notebook_runtime.NotebookRuntime: """Post-rpc interceptor for get_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime` interceptor runs + before the `post_get_notebook_runtime_with_metadata` interceptor. """ return response + def post_get_notebook_runtime_with_metadata( + self, + response: notebook_runtime.NotebookRuntime, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntime, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_with_metadata` interceptor runs after the + `post_get_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime` will be passed to + `post_get_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_get_notebook_runtime_template( self, request: notebook_service.GetNotebookRuntimeTemplateRequest, @@ -443,12 +632,38 @@ def post_get_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for get_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime_template` interceptor runs + before the `post_get_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_get_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_template_with_metadata` interceptor runs after the + `post_get_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime_template` will be passed to + `post_get_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_list_notebook_execution_jobs( self, request: notebook_service.ListNotebookExecutionJobsRequest, @@ -469,12 +684,38 @@ def post_list_notebook_execution_jobs( ) -> notebook_service.ListNotebookExecutionJobsResponse: """Post-rpc interceptor for list_notebook_execution_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_execution_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_execution_jobs` interceptor runs + before the `post_list_notebook_execution_jobs_with_metadata` interceptor. """ return response + def post_list_notebook_execution_jobs_with_metadata( + self, + response: notebook_service.ListNotebookExecutionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookExecutionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_execution_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_execution_jobs_with_metadata` + interceptor in new development instead of the `post_list_notebook_execution_jobs` interceptor. + When both interceptors are used, this `post_list_notebook_execution_jobs_with_metadata` interceptor runs after the + `post_list_notebook_execution_jobs` interceptor. The (possibly modified) response returned by + `post_list_notebook_execution_jobs` will be passed to + `post_list_notebook_execution_jobs_with_metadata`. + """ + return response, metadata + def pre_list_notebook_runtimes( self, request: notebook_service.ListNotebookRuntimesRequest, @@ -495,12 +736,38 @@ def post_list_notebook_runtimes( ) -> notebook_service.ListNotebookRuntimesResponse: """Post-rpc interceptor for list_notebook_runtimes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtimes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtimes` interceptor runs + before the `post_list_notebook_runtimes_with_metadata` interceptor. """ return response + def post_list_notebook_runtimes_with_metadata( + self, + response: notebook_service.ListNotebookRuntimesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtimes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtimes_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtimes` interceptor. + When both interceptors are used, this `post_list_notebook_runtimes_with_metadata` interceptor runs after the + `post_list_notebook_runtimes` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtimes` will be passed to + `post_list_notebook_runtimes_with_metadata`. + """ + return response, metadata + def pre_list_notebook_runtime_templates( self, request: notebook_service.ListNotebookRuntimeTemplatesRequest, @@ -521,12 +788,38 @@ def post_list_notebook_runtime_templates( ) -> notebook_service.ListNotebookRuntimeTemplatesResponse: """Post-rpc interceptor for list_notebook_runtime_templates - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtime_templates_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtime_templates` interceptor runs + before the `post_list_notebook_runtime_templates_with_metadata` interceptor. """ return response + def post_list_notebook_runtime_templates_with_metadata( + self, + response: notebook_service.ListNotebookRuntimeTemplatesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimeTemplatesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtime_templates + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtime_templates_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtime_templates` interceptor. + When both interceptors are used, this `post_list_notebook_runtime_templates_with_metadata` interceptor runs after the + `post_list_notebook_runtime_templates` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtime_templates` will be passed to + `post_list_notebook_runtime_templates_with_metadata`. + """ + return response, metadata + def pre_start_notebook_runtime( self, request: notebook_service.StartNotebookRuntimeRequest, @@ -547,12 +840,35 @@ def post_start_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for start_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_start_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_start_notebook_runtime` interceptor runs + before the `post_start_notebook_runtime_with_metadata` interceptor. """ return response + def post_start_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for start_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_start_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_start_notebook_runtime` interceptor. + When both interceptors are used, this `post_start_notebook_runtime_with_metadata` interceptor runs after the + `post_start_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_start_notebook_runtime` will be passed to + `post_start_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_stop_notebook_runtime( self, request: notebook_service.StopNotebookRuntimeRequest, @@ -573,12 +889,35 @@ def post_stop_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for stop_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_stop_notebook_runtime` interceptor runs + before the `post_stop_notebook_runtime_with_metadata` interceptor. """ return response + def post_stop_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_stop_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_stop_notebook_runtime` interceptor. + When both interceptors are used, this `post_stop_notebook_runtime_with_metadata` interceptor runs after the + `post_stop_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_stop_notebook_runtime` will be passed to + `post_stop_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_update_notebook_runtime_template( self, request: notebook_service.UpdateNotebookRuntimeTemplateRequest, @@ -599,12 +938,38 @@ def post_update_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for update_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_update_notebook_runtime_template` interceptor runs + before the `post_update_notebook_runtime_template_with_metadata` interceptor. """ return response + def post_update_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_update_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_update_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_update_notebook_runtime_template_with_metadata` interceptor runs after the + `post_update_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_update_notebook_runtime_template` will be passed to + `post_update_notebook_runtime_template_with_metadata`. + """ + return response, metadata + def pre_upgrade_notebook_runtime( self, request: notebook_service.UpgradeNotebookRuntimeRequest, @@ -625,12 +990,35 @@ def post_upgrade_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for upgrade_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upgrade_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_upgrade_notebook_runtime` interceptor runs + before the `post_upgrade_notebook_runtime_with_metadata` interceptor. """ return response + def post_upgrade_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upgrade_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_upgrade_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_upgrade_notebook_runtime` interceptor. + When both interceptors are used, this `post_upgrade_notebook_runtime_with_metadata` interceptor runs after the + `post_upgrade_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_upgrade_notebook_runtime` will be passed to + `post_upgrade_notebook_runtime_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3180,6 +3568,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_assign_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_assign_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3335,6 +3727,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3488,6 +3887,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3637,6 +4043,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3785,6 +4198,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3932,6 +4349,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4082,6 +4506,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4233,6 +4661,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4388,6 +4820,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_get_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4540,6 +4979,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_execution_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_notebook_execution_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4691,6 +5134,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_runtimes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_notebook_runtimes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4841,6 +5288,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_notebook_runtime_templates(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_list_notebook_runtime_templates_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4997,6 +5451,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_start_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_start_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5149,6 +5607,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_stop_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stop_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5306,6 +5768,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5462,6 +5931,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_upgrade_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upgrade_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest_asyncio.py index b04a5d760b..edca32d00f 100644 --- a/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/notebook_service/transports/rest_asyncio.py @@ -253,12 +253,35 @@ async def post_assign_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for assign_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_assign_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_assign_notebook_runtime` interceptor runs + before the `post_assign_notebook_runtime_with_metadata` interceptor. """ return response + async def post_assign_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for assign_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_assign_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_assign_notebook_runtime` interceptor. + When both interceptors are used, this `post_assign_notebook_runtime_with_metadata` interceptor runs after the + `post_assign_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_assign_notebook_runtime` will be passed to + `post_assign_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_create_notebook_execution_job( self, request: notebook_service.CreateNotebookExecutionJobRequest, @@ -279,12 +302,35 @@ async def post_create_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_execution_job` interceptor runs + before the `post_create_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_create_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_create_notebook_execution_job` interceptor. + When both interceptors are used, this `post_create_notebook_execution_job_with_metadata` interceptor runs after the + `post_create_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_create_notebook_execution_job` will be passed to + `post_create_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_create_notebook_runtime_template( self, request: notebook_service.CreateNotebookRuntimeTemplateRequest, @@ -305,12 +351,35 @@ async def post_create_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_create_notebook_runtime_template` interceptor runs + before the `post_create_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_create_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_create_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_create_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_create_notebook_runtime_template_with_metadata` interceptor runs after the + `post_create_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_create_notebook_runtime_template` will be passed to + `post_create_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_execution_job( self, request: notebook_service.DeleteNotebookExecutionJobRequest, @@ -331,12 +400,35 @@ async def post_delete_notebook_execution_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_execution_job` interceptor runs + before the `post_delete_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_delete_notebook_execution_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_delete_notebook_execution_job` interceptor. + When both interceptors are used, this `post_delete_notebook_execution_job_with_metadata` interceptor runs after the + `post_delete_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_delete_notebook_execution_job` will be passed to + `post_delete_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_runtime( self, request: notebook_service.DeleteNotebookRuntimeRequest, @@ -357,12 +449,35 @@ async def post_delete_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime` interceptor runs + before the `post_delete_notebook_runtime_with_metadata` interceptor. """ return response + async def post_delete_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_with_metadata` interceptor runs after the + `post_delete_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime` will be passed to + `post_delete_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_delete_notebook_runtime_template( self, request: notebook_service.DeleteNotebookRuntimeTemplateRequest, @@ -383,12 +498,35 @@ async def post_delete_notebook_runtime_template( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_delete_notebook_runtime_template` interceptor runs + before the `post_delete_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_delete_notebook_runtime_template_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_delete_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_delete_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_delete_notebook_runtime_template_with_metadata` interceptor runs after the + `post_delete_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_delete_notebook_runtime_template` will be passed to + `post_delete_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_execution_job( self, request: notebook_service.GetNotebookExecutionJobRequest, @@ -409,12 +547,38 @@ async def post_get_notebook_execution_job( ) -> notebook_execution_job.NotebookExecutionJob: """Post-rpc interceptor for get_notebook_execution_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_execution_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_execution_job` interceptor runs + before the `post_get_notebook_execution_job_with_metadata` interceptor. """ return response + async def post_get_notebook_execution_job_with_metadata( + self, + response: notebook_execution_job.NotebookExecutionJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_execution_job.NotebookExecutionJob, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_execution_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_execution_job_with_metadata` + interceptor in new development instead of the `post_get_notebook_execution_job` interceptor. + When both interceptors are used, this `post_get_notebook_execution_job_with_metadata` interceptor runs after the + `post_get_notebook_execution_job` interceptor. The (possibly modified) response returned by + `post_get_notebook_execution_job` will be passed to + `post_get_notebook_execution_job_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_runtime( self, request: notebook_service.GetNotebookRuntimeRequest, @@ -435,12 +599,37 @@ async def post_get_notebook_runtime( ) -> notebook_runtime.NotebookRuntime: """Post-rpc interceptor for get_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime` interceptor runs + before the `post_get_notebook_runtime_with_metadata` interceptor. """ return response + async def post_get_notebook_runtime_with_metadata( + self, + response: notebook_runtime.NotebookRuntime, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntime, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_with_metadata` interceptor runs after the + `post_get_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime` will be passed to + `post_get_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_get_notebook_runtime_template( self, request: notebook_service.GetNotebookRuntimeTemplateRequest, @@ -461,12 +650,38 @@ async def post_get_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for get_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_get_notebook_runtime_template` interceptor runs + before the `post_get_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_get_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_get_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_get_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_get_notebook_runtime_template_with_metadata` interceptor runs after the + `post_get_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_get_notebook_runtime_template` will be passed to + `post_get_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_execution_jobs( self, request: notebook_service.ListNotebookExecutionJobsRequest, @@ -487,12 +702,38 @@ async def post_list_notebook_execution_jobs( ) -> notebook_service.ListNotebookExecutionJobsResponse: """Post-rpc interceptor for list_notebook_execution_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_execution_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_execution_jobs` interceptor runs + before the `post_list_notebook_execution_jobs_with_metadata` interceptor. """ return response + async def post_list_notebook_execution_jobs_with_metadata( + self, + response: notebook_service.ListNotebookExecutionJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookExecutionJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_execution_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_execution_jobs_with_metadata` + interceptor in new development instead of the `post_list_notebook_execution_jobs` interceptor. + When both interceptors are used, this `post_list_notebook_execution_jobs_with_metadata` interceptor runs after the + `post_list_notebook_execution_jobs` interceptor. The (possibly modified) response returned by + `post_list_notebook_execution_jobs` will be passed to + `post_list_notebook_execution_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_runtimes( self, request: notebook_service.ListNotebookRuntimesRequest, @@ -513,12 +754,38 @@ async def post_list_notebook_runtimes( ) -> notebook_service.ListNotebookRuntimesResponse: """Post-rpc interceptor for list_notebook_runtimes - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtimes_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtimes` interceptor runs + before the `post_list_notebook_runtimes_with_metadata` interceptor. """ return response + async def post_list_notebook_runtimes_with_metadata( + self, + response: notebook_service.ListNotebookRuntimesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtimes + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtimes_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtimes` interceptor. + When both interceptors are used, this `post_list_notebook_runtimes_with_metadata` interceptor runs after the + `post_list_notebook_runtimes` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtimes` will be passed to + `post_list_notebook_runtimes_with_metadata`. + """ + return response, metadata + async def pre_list_notebook_runtime_templates( self, request: notebook_service.ListNotebookRuntimeTemplatesRequest, @@ -539,12 +806,38 @@ async def post_list_notebook_runtime_templates( ) -> notebook_service.ListNotebookRuntimeTemplatesResponse: """Post-rpc interceptor for list_notebook_runtime_templates - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_notebook_runtime_templates_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_list_notebook_runtime_templates` interceptor runs + before the `post_list_notebook_runtime_templates_with_metadata` interceptor. """ return response + async def post_list_notebook_runtime_templates_with_metadata( + self, + response: notebook_service.ListNotebookRuntimeTemplatesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_service.ListNotebookRuntimeTemplatesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_notebook_runtime_templates + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_list_notebook_runtime_templates_with_metadata` + interceptor in new development instead of the `post_list_notebook_runtime_templates` interceptor. + When both interceptors are used, this `post_list_notebook_runtime_templates_with_metadata` interceptor runs after the + `post_list_notebook_runtime_templates` interceptor. The (possibly modified) response returned by + `post_list_notebook_runtime_templates` will be passed to + `post_list_notebook_runtime_templates_with_metadata`. + """ + return response, metadata + async def pre_start_notebook_runtime( self, request: notebook_service.StartNotebookRuntimeRequest, @@ -565,12 +858,35 @@ async def post_start_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for start_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_start_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_start_notebook_runtime` interceptor runs + before the `post_start_notebook_runtime_with_metadata` interceptor. """ return response + async def post_start_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for start_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_start_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_start_notebook_runtime` interceptor. + When both interceptors are used, this `post_start_notebook_runtime_with_metadata` interceptor runs after the + `post_start_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_start_notebook_runtime` will be passed to + `post_start_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_stop_notebook_runtime( self, request: notebook_service.StopNotebookRuntimeRequest, @@ -591,12 +907,35 @@ async def post_stop_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for stop_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_stop_notebook_runtime` interceptor runs + before the `post_stop_notebook_runtime_with_metadata` interceptor. """ return response + async def post_stop_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_stop_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_stop_notebook_runtime` interceptor. + When both interceptors are used, this `post_stop_notebook_runtime_with_metadata` interceptor runs after the + `post_stop_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_stop_notebook_runtime` will be passed to + `post_stop_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_update_notebook_runtime_template( self, request: notebook_service.UpdateNotebookRuntimeTemplateRequest, @@ -617,12 +956,38 @@ async def post_update_notebook_runtime_template( ) -> notebook_runtime.NotebookRuntimeTemplate: """Post-rpc interceptor for update_notebook_runtime_template - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_notebook_runtime_template_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_update_notebook_runtime_template` interceptor runs + before the `post_update_notebook_runtime_template_with_metadata` interceptor. """ return response + async def post_update_notebook_runtime_template_with_metadata( + self, + response: notebook_runtime.NotebookRuntimeTemplate, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + notebook_runtime.NotebookRuntimeTemplate, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_notebook_runtime_template + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_update_notebook_runtime_template_with_metadata` + interceptor in new development instead of the `post_update_notebook_runtime_template` interceptor. + When both interceptors are used, this `post_update_notebook_runtime_template_with_metadata` interceptor runs after the + `post_update_notebook_runtime_template` interceptor. The (possibly modified) response returned by + `post_update_notebook_runtime_template` will be passed to + `post_update_notebook_runtime_template_with_metadata`. + """ + return response, metadata + async def pre_upgrade_notebook_runtime( self, request: notebook_service.UpgradeNotebookRuntimeRequest, @@ -643,12 +1008,35 @@ async def post_upgrade_notebook_runtime( ) -> operations_pb2.Operation: """Post-rpc interceptor for upgrade_notebook_runtime - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upgrade_notebook_runtime_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the NotebookService server but before - it is returned to user code. + it is returned to user code. This `post_upgrade_notebook_runtime` interceptor runs + before the `post_upgrade_notebook_runtime_with_metadata` interceptor. """ return response + async def post_upgrade_notebook_runtime_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for upgrade_notebook_runtime + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the NotebookService server but before it is returned to user code. + + We recommend only using this `post_upgrade_notebook_runtime_with_metadata` + interceptor in new development instead of the `post_upgrade_notebook_runtime` interceptor. + When both interceptors are used, this `post_upgrade_notebook_runtime_with_metadata` interceptor runs after the + `post_upgrade_notebook_runtime` interceptor. The (possibly modified) response returned by + `post_upgrade_notebook_runtime` will be passed to + `post_upgrade_notebook_runtime_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1238,6 +1626,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_assign_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_assign_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1402,6 +1797,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1568,6 +1970,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1726,6 +2135,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1880,6 +2296,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2040,6 +2463,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2194,6 +2624,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_execution_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_notebook_execution_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2351,6 +2788,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2513,6 +2954,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2672,6 +3120,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_execution_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_notebook_execution_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2829,6 +3284,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_runtimes(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_notebook_runtimes_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2990,6 +3449,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_notebook_runtime_templates(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_notebook_runtime_templates_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3154,6 +3620,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_start_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_start_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3314,6 +3784,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_stop_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stop_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3482,6 +3956,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_notebook_runtime_template(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_notebook_runtime_template_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3644,6 +4125,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upgrade_notebook_runtime(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_upgrade_notebook_runtime_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/client.py b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/client.py index 4d487729e1..89fe4af3b7 100644 --- a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -605,6 +607,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1703,16 +1732,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1758,16 +1791,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1930,16 +1967,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2051,16 +2092,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2173,16 +2218,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2233,16 +2282,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2288,16 +2341,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2343,16 +2400,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest.py index aa53e1a4d0..926fc766af 100644 --- a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest.py @@ -154,12 +154,35 @@ def post_create_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_create_persistent_resource` interceptor runs + before the `post_create_persistent_resource_with_metadata` interceptor. """ return response + def post_create_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_create_persistent_resource_with_metadata` + interceptor in new development instead of the `post_create_persistent_resource` interceptor. + When both interceptors are used, this `post_create_persistent_resource_with_metadata` interceptor runs after the + `post_create_persistent_resource` interceptor. The (possibly modified) response returned by + `post_create_persistent_resource` will be passed to + `post_create_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_delete_persistent_resource( self, request: persistent_resource_service.DeletePersistentResourceRequest, @@ -180,12 +203,35 @@ def post_delete_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_delete_persistent_resource` interceptor runs + before the `post_delete_persistent_resource_with_metadata` interceptor. """ return response + def post_delete_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_delete_persistent_resource_with_metadata` + interceptor in new development instead of the `post_delete_persistent_resource` interceptor. + When both interceptors are used, this `post_delete_persistent_resource_with_metadata` interceptor runs after the + `post_delete_persistent_resource` interceptor. The (possibly modified) response returned by + `post_delete_persistent_resource` will be passed to + `post_delete_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_get_persistent_resource( self, request: persistent_resource_service.GetPersistentResourceRequest, @@ -206,12 +252,37 @@ def post_get_persistent_resource( ) -> persistent_resource.PersistentResource: """Post-rpc interceptor for get_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_get_persistent_resource` interceptor runs + before the `post_get_persistent_resource_with_metadata` interceptor. """ return response + def post_get_persistent_resource_with_metadata( + self, + response: persistent_resource.PersistentResource, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource.PersistentResource, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_get_persistent_resource_with_metadata` + interceptor in new development instead of the `post_get_persistent_resource` interceptor. + When both interceptors are used, this `post_get_persistent_resource_with_metadata` interceptor runs after the + `post_get_persistent_resource` interceptor. The (possibly modified) response returned by + `post_get_persistent_resource` will be passed to + `post_get_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_list_persistent_resources( self, request: persistent_resource_service.ListPersistentResourcesRequest, @@ -232,12 +303,38 @@ def post_list_persistent_resources( ) -> persistent_resource_service.ListPersistentResourcesResponse: """Post-rpc interceptor for list_persistent_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_persistent_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_list_persistent_resources` interceptor runs + before the `post_list_persistent_resources_with_metadata` interceptor. """ return response + def post_list_persistent_resources_with_metadata( + self, + response: persistent_resource_service.ListPersistentResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource_service.ListPersistentResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_persistent_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_list_persistent_resources_with_metadata` + interceptor in new development instead of the `post_list_persistent_resources` interceptor. + When both interceptors are used, this `post_list_persistent_resources_with_metadata` interceptor runs after the + `post_list_persistent_resources` interceptor. The (possibly modified) response returned by + `post_list_persistent_resources` will be passed to + `post_list_persistent_resources_with_metadata`. + """ + return response, metadata + def pre_reboot_persistent_resource( self, request: persistent_resource_service.RebootPersistentResourceRequest, @@ -258,12 +355,35 @@ def post_reboot_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for reboot_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_reboot_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_reboot_persistent_resource` interceptor runs + before the `post_reboot_persistent_resource_with_metadata` interceptor. """ return response + def post_reboot_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for reboot_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_reboot_persistent_resource_with_metadata` + interceptor in new development instead of the `post_reboot_persistent_resource` interceptor. + When both interceptors are used, this `post_reboot_persistent_resource_with_metadata` interceptor runs after the + `post_reboot_persistent_resource` interceptor. The (possibly modified) response returned by + `post_reboot_persistent_resource` will be passed to + `post_reboot_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_update_persistent_resource( self, request: persistent_resource_service.UpdatePersistentResourceRequest, @@ -284,12 +404,35 @@ def post_update_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_update_persistent_resource` interceptor runs + before the `post_update_persistent_resource_with_metadata` interceptor. """ return response + def post_update_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_update_persistent_resource_with_metadata` + interceptor in new development instead of the `post_update_persistent_resource` interceptor. + When both interceptors are used, this `post_update_persistent_resource_with_metadata` interceptor runs after the + `post_update_persistent_resource` interceptor. The (possibly modified) response returned by + `post_update_persistent_resource` will be passed to + `post_update_persistent_resource_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2842,6 +2985,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2991,6 +3138,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3141,6 +3292,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3294,6 +3449,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_persistent_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_persistent_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3451,6 +3610,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_reboot_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_reboot_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3606,6 +3769,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest_asyncio.py index a6f0f6162b..cb001d4e61 100644 --- a/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/persistent_resource_service/transports/rest_asyncio.py @@ -172,12 +172,35 @@ async def post_create_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_create_persistent_resource` interceptor runs + before the `post_create_persistent_resource_with_metadata` interceptor. """ return response + async def post_create_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_create_persistent_resource_with_metadata` + interceptor in new development instead of the `post_create_persistent_resource` interceptor. + When both interceptors are used, this `post_create_persistent_resource_with_metadata` interceptor runs after the + `post_create_persistent_resource` interceptor. The (possibly modified) response returned by + `post_create_persistent_resource` will be passed to + `post_create_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_delete_persistent_resource( self, request: persistent_resource_service.DeletePersistentResourceRequest, @@ -198,12 +221,35 @@ async def post_delete_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_delete_persistent_resource` interceptor runs + before the `post_delete_persistent_resource_with_metadata` interceptor. """ return response + async def post_delete_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_delete_persistent_resource_with_metadata` + interceptor in new development instead of the `post_delete_persistent_resource` interceptor. + When both interceptors are used, this `post_delete_persistent_resource_with_metadata` interceptor runs after the + `post_delete_persistent_resource` interceptor. The (possibly modified) response returned by + `post_delete_persistent_resource` will be passed to + `post_delete_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_get_persistent_resource( self, request: persistent_resource_service.GetPersistentResourceRequest, @@ -224,12 +270,37 @@ async def post_get_persistent_resource( ) -> persistent_resource.PersistentResource: """Post-rpc interceptor for get_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_get_persistent_resource` interceptor runs + before the `post_get_persistent_resource_with_metadata` interceptor. """ return response + async def post_get_persistent_resource_with_metadata( + self, + response: persistent_resource.PersistentResource, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource.PersistentResource, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_get_persistent_resource_with_metadata` + interceptor in new development instead of the `post_get_persistent_resource` interceptor. + When both interceptors are used, this `post_get_persistent_resource_with_metadata` interceptor runs after the + `post_get_persistent_resource` interceptor. The (possibly modified) response returned by + `post_get_persistent_resource` will be passed to + `post_get_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_list_persistent_resources( self, request: persistent_resource_service.ListPersistentResourcesRequest, @@ -250,12 +321,38 @@ async def post_list_persistent_resources( ) -> persistent_resource_service.ListPersistentResourcesResponse: """Post-rpc interceptor for list_persistent_resources - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_persistent_resources_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_list_persistent_resources` interceptor runs + before the `post_list_persistent_resources_with_metadata` interceptor. """ return response + async def post_list_persistent_resources_with_metadata( + self, + response: persistent_resource_service.ListPersistentResourcesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + persistent_resource_service.ListPersistentResourcesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_persistent_resources + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_list_persistent_resources_with_metadata` + interceptor in new development instead of the `post_list_persistent_resources` interceptor. + When both interceptors are used, this `post_list_persistent_resources_with_metadata` interceptor runs after the + `post_list_persistent_resources` interceptor. The (possibly modified) response returned by + `post_list_persistent_resources` will be passed to + `post_list_persistent_resources_with_metadata`. + """ + return response, metadata + async def pre_reboot_persistent_resource( self, request: persistent_resource_service.RebootPersistentResourceRequest, @@ -276,12 +373,35 @@ async def post_reboot_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for reboot_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_reboot_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_reboot_persistent_resource` interceptor runs + before the `post_reboot_persistent_resource_with_metadata` interceptor. """ return response + async def post_reboot_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for reboot_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_reboot_persistent_resource_with_metadata` + interceptor in new development instead of the `post_reboot_persistent_resource` interceptor. + When both interceptors are used, this `post_reboot_persistent_resource_with_metadata` interceptor runs after the + `post_reboot_persistent_resource` interceptor. The (possibly modified) response returned by + `post_reboot_persistent_resource` will be passed to + `post_reboot_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_update_persistent_resource( self, request: persistent_resource_service.UpdatePersistentResourceRequest, @@ -302,12 +422,35 @@ async def post_update_persistent_resource( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_persistent_resource - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_persistent_resource_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PersistentResourceService server but before - it is returned to user code. + it is returned to user code. This `post_update_persistent_resource` interceptor runs + before the `post_update_persistent_resource_with_metadata` interceptor. """ return response + async def post_update_persistent_resource_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_persistent_resource + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PersistentResourceService server but before it is returned to user code. + + We recommend only using this `post_update_persistent_resource_with_metadata` + interceptor in new development instead of the `post_update_persistent_resource` interceptor. + When both interceptors are used, this `post_update_persistent_resource_with_metadata` interceptor runs after the + `post_update_persistent_resource` interceptor. The (possibly modified) response returned by + `post_update_persistent_resource` will be passed to + `post_update_persistent_resource_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -854,6 +997,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1011,6 +1161,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1169,6 +1326,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1328,6 +1492,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_persistent_resources(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_persistent_resources_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1493,6 +1664,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_reboot_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_reboot_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1656,6 +1834,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_persistent_resource(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_persistent_resource_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/pipeline_service/client.py b/google/cloud/aiplatform_v1beta1/services/pipeline_service/client.py index 7c9d5c8e13..b0bb21babb 100644 --- a/google/cloud/aiplatform_v1beta1/services/pipeline_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/pipeline_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -723,6 +725,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2482,16 +2511,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2537,16 +2570,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2709,16 +2746,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2830,16 +2871,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2952,16 +2997,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3012,16 +3061,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3067,16 +3120,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3122,16 +3179,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest.py index a50dfe8d5b..4901a47a98 100644 --- a/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest.py @@ -200,12 +200,35 @@ def post_batch_cancel_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_cancel_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_cancel_pipeline_jobs` interceptor runs + before the `post_batch_cancel_pipeline_jobs_with_metadata` interceptor. """ return response + def post_batch_cancel_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_cancel_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_cancel_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_cancel_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_cancel_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_cancel_pipeline_jobs` will be passed to + `post_batch_cancel_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_batch_delete_pipeline_jobs( self, request: pipeline_service.BatchDeletePipelineJobsRequest, @@ -226,12 +249,35 @@ def post_batch_delete_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_delete_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_delete_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_delete_pipeline_jobs` interceptor runs + before the `post_batch_delete_pipeline_jobs_with_metadata` interceptor. """ return response + def post_batch_delete_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_delete_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_delete_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_delete_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_delete_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_delete_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_delete_pipeline_jobs` will be passed to + `post_batch_delete_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_cancel_pipeline_job( self, request: pipeline_service.CancelPipelineJobRequest, @@ -282,12 +328,35 @@ def post_create_pipeline_job( ) -> gca_pipeline_job.PipelineJob: """Post-rpc interceptor for create_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_pipeline_job` interceptor runs + before the `post_create_pipeline_job_with_metadata` interceptor. """ return response + def post_create_pipeline_job_with_metadata( + self, + response: gca_pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_pipeline_job_with_metadata` + interceptor in new development instead of the `post_create_pipeline_job` interceptor. + When both interceptors are used, this `post_create_pipeline_job_with_metadata` interceptor runs after the + `post_create_pipeline_job` interceptor. The (possibly modified) response returned by + `post_create_pipeline_job` will be passed to + `post_create_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_create_training_pipeline( self, request: pipeline_service.CreateTrainingPipelineRequest, @@ -308,12 +377,37 @@ def post_create_training_pipeline( ) -> gca_training_pipeline.TrainingPipeline: """Post-rpc interceptor for create_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_training_pipeline` interceptor runs + before the `post_create_training_pipeline_with_metadata` interceptor. """ return response + def post_create_training_pipeline_with_metadata( + self, + response: gca_training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_training_pipeline_with_metadata` + interceptor in new development instead of the `post_create_training_pipeline` interceptor. + When both interceptors are used, this `post_create_training_pipeline_with_metadata` interceptor runs after the + `post_create_training_pipeline` interceptor. The (possibly modified) response returned by + `post_create_training_pipeline` will be passed to + `post_create_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_delete_pipeline_job( self, request: pipeline_service.DeletePipelineJobRequest, @@ -334,12 +428,35 @@ def post_delete_pipeline_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_pipeline_job` interceptor runs + before the `post_delete_pipeline_job_with_metadata` interceptor. """ return response + def post_delete_pipeline_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_pipeline_job_with_metadata` + interceptor in new development instead of the `post_delete_pipeline_job` interceptor. + When both interceptors are used, this `post_delete_pipeline_job_with_metadata` interceptor runs after the + `post_delete_pipeline_job` interceptor. The (possibly modified) response returned by + `post_delete_pipeline_job` will be passed to + `post_delete_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_delete_training_pipeline( self, request: pipeline_service.DeleteTrainingPipelineRequest, @@ -360,12 +477,35 @@ def post_delete_training_pipeline( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_training_pipeline` interceptor runs + before the `post_delete_training_pipeline_with_metadata` interceptor. """ return response + def post_delete_training_pipeline_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_training_pipeline_with_metadata` + interceptor in new development instead of the `post_delete_training_pipeline` interceptor. + When both interceptors are used, this `post_delete_training_pipeline_with_metadata` interceptor runs after the + `post_delete_training_pipeline` interceptor. The (possibly modified) response returned by + `post_delete_training_pipeline` will be passed to + `post_delete_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_get_pipeline_job( self, request: pipeline_service.GetPipelineJobRequest, @@ -385,12 +525,35 @@ def post_get_pipeline_job( ) -> pipeline_job.PipelineJob: """Post-rpc interceptor for get_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_pipeline_job` interceptor runs + before the `post_get_pipeline_job_with_metadata` interceptor. """ return response + def post_get_pipeline_job_with_metadata( + self, + response: pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_pipeline_job_with_metadata` + interceptor in new development instead of the `post_get_pipeline_job` interceptor. + When both interceptors are used, this `post_get_pipeline_job_with_metadata` interceptor runs after the + `post_get_pipeline_job` interceptor. The (possibly modified) response returned by + `post_get_pipeline_job` will be passed to + `post_get_pipeline_job_with_metadata`. + """ + return response, metadata + def pre_get_training_pipeline( self, request: pipeline_service.GetTrainingPipelineRequest, @@ -411,12 +574,37 @@ def post_get_training_pipeline( ) -> training_pipeline.TrainingPipeline: """Post-rpc interceptor for get_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_training_pipeline` interceptor runs + before the `post_get_training_pipeline_with_metadata` interceptor. """ return response + def post_get_training_pipeline_with_metadata( + self, + response: training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_training_pipeline_with_metadata` + interceptor in new development instead of the `post_get_training_pipeline` interceptor. + When both interceptors are used, this `post_get_training_pipeline_with_metadata` interceptor runs after the + `post_get_training_pipeline` interceptor. The (possibly modified) response returned by + `post_get_training_pipeline` will be passed to + `post_get_training_pipeline_with_metadata`. + """ + return response, metadata + def pre_list_pipeline_jobs( self, request: pipeline_service.ListPipelineJobsRequest, @@ -437,12 +625,38 @@ def post_list_pipeline_jobs( ) -> pipeline_service.ListPipelineJobsResponse: """Post-rpc interceptor for list_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_pipeline_jobs` interceptor runs + before the `post_list_pipeline_jobs_with_metadata` interceptor. """ return response + def post_list_pipeline_jobs_with_metadata( + self, + response: pipeline_service.ListPipelineJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListPipelineJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_list_pipeline_jobs` interceptor. + When both interceptors are used, this `post_list_pipeline_jobs_with_metadata` interceptor runs after the + `post_list_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_list_pipeline_jobs` will be passed to + `post_list_pipeline_jobs_with_metadata`. + """ + return response, metadata + def pre_list_training_pipelines( self, request: pipeline_service.ListTrainingPipelinesRequest, @@ -463,12 +677,38 @@ def post_list_training_pipelines( ) -> pipeline_service.ListTrainingPipelinesResponse: """Post-rpc interceptor for list_training_pipelines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_training_pipelines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_training_pipelines` interceptor runs + before the `post_list_training_pipelines_with_metadata` interceptor. """ return response + def post_list_training_pipelines_with_metadata( + self, + response: pipeline_service.ListTrainingPipelinesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListTrainingPipelinesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_training_pipelines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_training_pipelines_with_metadata` + interceptor in new development instead of the `post_list_training_pipelines` interceptor. + When both interceptors are used, this `post_list_training_pipelines_with_metadata` interceptor runs after the + `post_list_training_pipelines` interceptor. The (possibly modified) response returned by + `post_list_training_pipelines` will be passed to + `post_list_training_pipelines_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3021,6 +3261,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_cancel_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_cancel_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3176,6 +3420,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_delete_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_batch_delete_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3565,6 +3813,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3724,6 +3976,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3872,6 +4128,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4020,6 +4280,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4166,6 +4430,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4317,6 +4585,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4465,6 +4737,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4616,6 +4892,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_training_pipelines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_training_pipelines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest_asyncio.py index ff2b16cdf6..5ab845879f 100644 --- a/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/pipeline_service/transports/rest_asyncio.py @@ -218,12 +218,35 @@ async def post_batch_cancel_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_cancel_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_cancel_pipeline_jobs` interceptor runs + before the `post_batch_cancel_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_batch_cancel_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_cancel_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_cancel_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_cancel_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_cancel_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_cancel_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_cancel_pipeline_jobs` will be passed to + `post_batch_cancel_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_batch_delete_pipeline_jobs( self, request: pipeline_service.BatchDeletePipelineJobsRequest, @@ -244,12 +267,35 @@ async def post_batch_delete_pipeline_jobs( ) -> operations_pb2.Operation: """Post-rpc interceptor for batch_delete_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_delete_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_batch_delete_pipeline_jobs` interceptor runs + before the `post_batch_delete_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_batch_delete_pipeline_jobs_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for batch_delete_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_batch_delete_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_batch_delete_pipeline_jobs` interceptor. + When both interceptors are used, this `post_batch_delete_pipeline_jobs_with_metadata` interceptor runs after the + `post_batch_delete_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_batch_delete_pipeline_jobs` will be passed to + `post_batch_delete_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_cancel_pipeline_job( self, request: pipeline_service.CancelPipelineJobRequest, @@ -300,12 +346,35 @@ async def post_create_pipeline_job( ) -> gca_pipeline_job.PipelineJob: """Post-rpc interceptor for create_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_pipeline_job` interceptor runs + before the `post_create_pipeline_job_with_metadata` interceptor. """ return response + async def post_create_pipeline_job_with_metadata( + self, + response: gca_pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_pipeline_job_with_metadata` + interceptor in new development instead of the `post_create_pipeline_job` interceptor. + When both interceptors are used, this `post_create_pipeline_job_with_metadata` interceptor runs after the + `post_create_pipeline_job` interceptor. The (possibly modified) response returned by + `post_create_pipeline_job` will be passed to + `post_create_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_create_training_pipeline( self, request: pipeline_service.CreateTrainingPipelineRequest, @@ -326,12 +395,37 @@ async def post_create_training_pipeline( ) -> gca_training_pipeline.TrainingPipeline: """Post-rpc interceptor for create_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_create_training_pipeline` interceptor runs + before the `post_create_training_pipeline_with_metadata` interceptor. """ return response + async def post_create_training_pipeline_with_metadata( + self, + response: gca_training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_create_training_pipeline_with_metadata` + interceptor in new development instead of the `post_create_training_pipeline` interceptor. + When both interceptors are used, this `post_create_training_pipeline_with_metadata` interceptor runs after the + `post_create_training_pipeline` interceptor. The (possibly modified) response returned by + `post_create_training_pipeline` will be passed to + `post_create_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_delete_pipeline_job( self, request: pipeline_service.DeletePipelineJobRequest, @@ -352,12 +446,35 @@ async def post_delete_pipeline_job( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_pipeline_job` interceptor runs + before the `post_delete_pipeline_job_with_metadata` interceptor. """ return response + async def post_delete_pipeline_job_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_pipeline_job_with_metadata` + interceptor in new development instead of the `post_delete_pipeline_job` interceptor. + When both interceptors are used, this `post_delete_pipeline_job_with_metadata` interceptor runs after the + `post_delete_pipeline_job` interceptor. The (possibly modified) response returned by + `post_delete_pipeline_job` will be passed to + `post_delete_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_delete_training_pipeline( self, request: pipeline_service.DeleteTrainingPipelineRequest, @@ -378,12 +495,35 @@ async def post_delete_training_pipeline( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_training_pipeline` interceptor runs + before the `post_delete_training_pipeline_with_metadata` interceptor. """ return response + async def post_delete_training_pipeline_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_delete_training_pipeline_with_metadata` + interceptor in new development instead of the `post_delete_training_pipeline` interceptor. + When both interceptors are used, this `post_delete_training_pipeline_with_metadata` interceptor runs after the + `post_delete_training_pipeline` interceptor. The (possibly modified) response returned by + `post_delete_training_pipeline` will be passed to + `post_delete_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_get_pipeline_job( self, request: pipeline_service.GetPipelineJobRequest, @@ -403,12 +543,35 @@ async def post_get_pipeline_job( ) -> pipeline_job.PipelineJob: """Post-rpc interceptor for get_pipeline_job - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_pipeline_job_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_pipeline_job` interceptor runs + before the `post_get_pipeline_job_with_metadata` interceptor. """ return response + async def post_get_pipeline_job_with_metadata( + self, + response: pipeline_job.PipelineJob, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[pipeline_job.PipelineJob, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_pipeline_job + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_pipeline_job_with_metadata` + interceptor in new development instead of the `post_get_pipeline_job` interceptor. + When both interceptors are used, this `post_get_pipeline_job_with_metadata` interceptor runs after the + `post_get_pipeline_job` interceptor. The (possibly modified) response returned by + `post_get_pipeline_job` will be passed to + `post_get_pipeline_job_with_metadata`. + """ + return response, metadata + async def pre_get_training_pipeline( self, request: pipeline_service.GetTrainingPipelineRequest, @@ -429,12 +592,37 @@ async def post_get_training_pipeline( ) -> training_pipeline.TrainingPipeline: """Post-rpc interceptor for get_training_pipeline - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_training_pipeline_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_get_training_pipeline` interceptor runs + before the `post_get_training_pipeline_with_metadata` interceptor. """ return response + async def post_get_training_pipeline_with_metadata( + self, + response: training_pipeline.TrainingPipeline, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + training_pipeline.TrainingPipeline, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_training_pipeline + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_get_training_pipeline_with_metadata` + interceptor in new development instead of the `post_get_training_pipeline` interceptor. + When both interceptors are used, this `post_get_training_pipeline_with_metadata` interceptor runs after the + `post_get_training_pipeline` interceptor. The (possibly modified) response returned by + `post_get_training_pipeline` will be passed to + `post_get_training_pipeline_with_metadata`. + """ + return response, metadata + async def pre_list_pipeline_jobs( self, request: pipeline_service.ListPipelineJobsRequest, @@ -455,12 +643,38 @@ async def post_list_pipeline_jobs( ) -> pipeline_service.ListPipelineJobsResponse: """Post-rpc interceptor for list_pipeline_jobs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_pipeline_jobs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_pipeline_jobs` interceptor runs + before the `post_list_pipeline_jobs_with_metadata` interceptor. """ return response + async def post_list_pipeline_jobs_with_metadata( + self, + response: pipeline_service.ListPipelineJobsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListPipelineJobsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_pipeline_jobs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_pipeline_jobs_with_metadata` + interceptor in new development instead of the `post_list_pipeline_jobs` interceptor. + When both interceptors are used, this `post_list_pipeline_jobs_with_metadata` interceptor runs after the + `post_list_pipeline_jobs` interceptor. The (possibly modified) response returned by + `post_list_pipeline_jobs` will be passed to + `post_list_pipeline_jobs_with_metadata`. + """ + return response, metadata + async def pre_list_training_pipelines( self, request: pipeline_service.ListTrainingPipelinesRequest, @@ -481,12 +695,38 @@ async def post_list_training_pipelines( ) -> pipeline_service.ListTrainingPipelinesResponse: """Post-rpc interceptor for list_training_pipelines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_training_pipelines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PipelineService server but before - it is returned to user code. + it is returned to user code. This `post_list_training_pipelines` interceptor runs + before the `post_list_training_pipelines_with_metadata` interceptor. """ return response + async def post_list_training_pipelines_with_metadata( + self, + response: pipeline_service.ListTrainingPipelinesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + pipeline_service.ListTrainingPipelinesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_training_pipelines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PipelineService server but before it is returned to user code. + + We recommend only using this `post_list_training_pipelines_with_metadata` + interceptor in new development instead of the `post_list_training_pipelines` interceptor. + When both interceptors are used, this `post_list_training_pipelines_with_metadata` interceptor runs after the + `post_list_training_pipelines` interceptor. The (possibly modified) response returned by + `post_list_training_pipelines` will be passed to + `post_list_training_pipelines_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1059,6 +1299,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_cancel_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_cancel_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1220,6 +1467,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_delete_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_delete_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1625,6 +1879,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1788,6 +2046,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1944,6 +2209,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2098,6 +2367,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2253,6 +2529,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_pipeline_job(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_pipeline_job_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2410,6 +2690,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_training_pipeline(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_training_pipeline_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2567,6 +2851,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_pipeline_jobs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_pipeline_jobs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2722,6 +3010,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_training_pipelines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_training_pipelines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/prediction_service/async_client.py b/google/cloud/aiplatform_v1beta1/services/prediction_service/async_client.py index f76710287b..89eb6227ef 100644 --- a/google/cloud/aiplatform_v1beta1/services/prediction_service/async_client.py +++ b/google/cloud/aiplatform_v1beta1/services/prediction_service/async_client.py @@ -58,6 +58,7 @@ from google.longrunning import operations_pb2 # type: ignore from google.protobuf import any_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from .transports.base import PredictionServiceTransport, DEFAULT_CLIENT_INFO from .transports.grpc_asyncio import PredictionServiceGrpcAsyncIOTransport from .client import PredictionServiceClient diff --git a/google/cloud/aiplatform_v1beta1/services/prediction_service/client.py b/google/cloud/aiplatform_v1beta1/services/prediction_service/client.py index 91fd76407c..a30173b7a0 100644 --- a/google/cloud/aiplatform_v1beta1/services/prediction_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/prediction_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -71,6 +73,7 @@ from google.longrunning import operations_pb2 # type: ignore from google.protobuf import any_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from .transports.base import PredictionServiceTransport, DEFAULT_CLIENT_INFO from .transports.grpc import PredictionServiceGrpcTransport from .transports.grpc_asyncio import PredictionServiceGrpcAsyncIOTransport @@ -574,6 +577,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2708,16 +2738,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2763,16 +2797,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2935,16 +2973,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -3056,16 +3098,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -3178,16 +3224,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -3238,16 +3288,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -3293,16 +3347,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3348,16 +3406,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest.py index 7e28e266de..f7459494f0 100644 --- a/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest.py @@ -193,12 +193,37 @@ def post_chat_completions( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for chat_completions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_chat_completions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_chat_completions` interceptor runs + before the `post_chat_completions_with_metadata` interceptor. """ return response + def post_chat_completions_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for chat_completions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_chat_completions_with_metadata` + interceptor in new development instead of the `post_chat_completions` interceptor. + When both interceptors are used, this `post_chat_completions_with_metadata` interceptor runs after the + `post_chat_completions` interceptor. The (possibly modified) response returned by + `post_chat_completions` will be passed to + `post_chat_completions_with_metadata`. + """ + return response, metadata + def pre_count_tokens( self, request: prediction_service.CountTokensRequest, @@ -218,12 +243,37 @@ def post_count_tokens( ) -> prediction_service.CountTokensResponse: """Post-rpc interceptor for count_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_count_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_count_tokens` interceptor runs + before the `post_count_tokens_with_metadata` interceptor. """ return response + def post_count_tokens_with_metadata( + self, + response: prediction_service.CountTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.CountTokensResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for count_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_count_tokens_with_metadata` + interceptor in new development instead of the `post_count_tokens` interceptor. + When both interceptors are used, this `post_count_tokens_with_metadata` interceptor runs after the + `post_count_tokens` interceptor. The (possibly modified) response returned by + `post_count_tokens` will be passed to + `post_count_tokens_with_metadata`. + """ + return response, metadata + def pre_direct_predict( self, request: prediction_service.DirectPredictRequest, @@ -243,12 +293,38 @@ def post_direct_predict( ) -> prediction_service.DirectPredictResponse: """Post-rpc interceptor for direct_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_predict` interceptor runs + before the `post_direct_predict_with_metadata` interceptor. """ return response + def post_direct_predict_with_metadata( + self, + response: prediction_service.DirectPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_predict_with_metadata` + interceptor in new development instead of the `post_direct_predict` interceptor. + When both interceptors are used, this `post_direct_predict_with_metadata` interceptor runs after the + `post_direct_predict` interceptor. The (possibly modified) response returned by + `post_direct_predict` will be passed to + `post_direct_predict_with_metadata`. + """ + return response, metadata + def pre_direct_raw_predict( self, request: prediction_service.DirectRawPredictRequest, @@ -269,12 +345,38 @@ def post_direct_raw_predict( ) -> prediction_service.DirectRawPredictResponse: """Post-rpc interceptor for direct_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_raw_predict` interceptor runs + before the `post_direct_raw_predict_with_metadata` interceptor. """ return response + def post_direct_raw_predict_with_metadata( + self, + response: prediction_service.DirectRawPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectRawPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_raw_predict_with_metadata` + interceptor in new development instead of the `post_direct_raw_predict` interceptor. + When both interceptors are used, this `post_direct_raw_predict_with_metadata` interceptor runs after the + `post_direct_raw_predict` interceptor. The (possibly modified) response returned by + `post_direct_raw_predict` will be passed to + `post_direct_raw_predict_with_metadata`. + """ + return response, metadata + def pre_explain( self, request: prediction_service.ExplainRequest, @@ -294,12 +396,37 @@ def post_explain( ) -> prediction_service.ExplainResponse: """Post-rpc interceptor for explain - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_explain_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_explain` interceptor runs + before the `post_explain_with_metadata` interceptor. """ return response + def post_explain_with_metadata( + self, + response: prediction_service.ExplainResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.ExplainResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for explain + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_explain_with_metadata` + interceptor in new development instead of the `post_explain` interceptor. + When both interceptors are used, this `post_explain_with_metadata` interceptor runs after the + `post_explain` interceptor. The (possibly modified) response returned by + `post_explain` will be passed to + `post_explain_with_metadata`. + """ + return response, metadata + def pre_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -320,12 +447,38 @@ def post_generate_content( ) -> prediction_service.GenerateContentResponse: """Post-rpc interceptor for generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_generate_content` interceptor runs + before the `post_generate_content_with_metadata` interceptor. """ return response + def post_generate_content_with_metadata( + self, + response: prediction_service.GenerateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.GenerateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_generate_content_with_metadata` + interceptor in new development instead of the `post_generate_content` interceptor. + When both interceptors are used, this `post_generate_content_with_metadata` interceptor runs after the + `post_generate_content` interceptor. The (possibly modified) response returned by + `post_generate_content` will be passed to + `post_generate_content_with_metadata`. + """ + return response, metadata + def pre_predict( self, request: prediction_service.PredictRequest, @@ -345,12 +498,37 @@ def post_predict( ) -> prediction_service.PredictResponse: """Post-rpc interceptor for predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_predict` interceptor runs + before the `post_predict_with_metadata` interceptor. """ return response + def post_predict_with_metadata( + self, + response: prediction_service.PredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.PredictResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_predict_with_metadata` + interceptor in new development instead of the `post_predict` interceptor. + When both interceptors are used, this `post_predict_with_metadata` interceptor runs after the + `post_predict` interceptor. The (possibly modified) response returned by + `post_predict` will be passed to + `post_predict_with_metadata`. + """ + return response, metadata + def pre_raw_predict( self, request: prediction_service.RawPredictRequest, @@ -370,12 +548,35 @@ def post_raw_predict( ) -> httpbody_pb2.HttpBody: """Post-rpc interceptor for raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_raw_predict` interceptor runs + before the `post_raw_predict_with_metadata` interceptor. """ return response + def post_raw_predict_with_metadata( + self, + response: httpbody_pb2.HttpBody, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[httpbody_pb2.HttpBody, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_raw_predict_with_metadata` + interceptor in new development instead of the `post_raw_predict` interceptor. + When both interceptors are used, this `post_raw_predict_with_metadata` interceptor runs after the + `post_raw_predict` interceptor. The (possibly modified) response returned by + `post_raw_predict` will be passed to + `post_raw_predict_with_metadata`. + """ + return response, metadata + def pre_server_streaming_predict( self, request: prediction_service.StreamingPredictRequest, @@ -396,12 +597,37 @@ def post_server_streaming_predict( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for server_streaming_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_server_streaming_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_server_streaming_predict` interceptor runs + before the `post_server_streaming_predict_with_metadata` interceptor. """ return response + def post_server_streaming_predict_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for server_streaming_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_server_streaming_predict_with_metadata` + interceptor in new development instead of the `post_server_streaming_predict` interceptor. + When both interceptors are used, this `post_server_streaming_predict_with_metadata` interceptor runs after the + `post_server_streaming_predict` interceptor. The (possibly modified) response returned by + `post_server_streaming_predict` will be passed to + `post_server_streaming_predict_with_metadata`. + """ + return response, metadata + def pre_stream_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -422,12 +648,37 @@ def post_stream_generate_content( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_generate_content` interceptor runs + before the `post_stream_generate_content_with_metadata` interceptor. """ return response + def post_stream_generate_content_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_generate_content_with_metadata` + interceptor in new development instead of the `post_stream_generate_content` interceptor. + When both interceptors are used, this `post_stream_generate_content_with_metadata` interceptor runs after the + `post_stream_generate_content` interceptor. The (possibly modified) response returned by + `post_stream_generate_content` will be passed to + `post_stream_generate_content_with_metadata`. + """ + return response, metadata + def pre_stream_raw_predict( self, request: prediction_service.StreamRawPredictRequest, @@ -448,12 +699,37 @@ def post_stream_raw_predict( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_raw_predict` interceptor runs + before the `post_stream_raw_predict_with_metadata` interceptor. """ return response + def post_stream_raw_predict_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_raw_predict_with_metadata` + interceptor in new development instead of the `post_stream_raw_predict` interceptor. + When both interceptors are used, this `post_stream_raw_predict_with_metadata` interceptor runs after the + `post_stream_raw_predict` interceptor. The (possibly modified) response returned by + `post_stream_raw_predict` will be passed to + `post_stream_raw_predict_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -958,6 +1234,10 @@ def __call__( resp = rest_streaming.ResponseIterator(response, httpbody_pb2.HttpBody) resp = self._interceptor.post_chat_completions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_chat_completions_with_metadata( + resp, response_metadata + ) return resp class _CountTokens( @@ -1087,6 +1367,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_count_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_count_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1240,6 +1524,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_direct_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_direct_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1395,6 +1683,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_direct_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_direct_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1551,6 +1843,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_explain(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_explain_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1705,6 +2001,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_generate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1861,6 +2161,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2060,6 +2364,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2220,6 +2528,10 @@ def __call__( ) resp = self._interceptor.post_server_streaming_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_server_streaming_predict_with_metadata( + resp, response_metadata + ) return resp class _StreamDirectPredict( @@ -2391,6 +2703,10 @@ def __call__( ) resp = self._interceptor.post_stream_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stream_generate_content_with_metadata( + resp, response_metadata + ) return resp class _StreamingPredict( @@ -2606,6 +2922,10 @@ def __call__( resp = rest_streaming.ResponseIterator(response, httpbody_pb2.HttpBody) resp = self._interceptor.post_stream_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stream_raw_predict_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest_asyncio.py index de906c89e0..bd5f3bfe8b 100644 --- a/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/prediction_service/transports/rest_asyncio.py @@ -210,12 +210,38 @@ async def post_chat_completions( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for chat_completions - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_chat_completions_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_chat_completions` interceptor runs + before the `post_chat_completions_with_metadata` interceptor. """ return response + async def post_chat_completions_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for chat_completions + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_chat_completions_with_metadata` + interceptor in new development instead of the `post_chat_completions` interceptor. + When both interceptors are used, this `post_chat_completions_with_metadata` interceptor runs after the + `post_chat_completions` interceptor. The (possibly modified) response returned by + `post_chat_completions` will be passed to + `post_chat_completions_with_metadata`. + """ + return response, metadata + async def pre_count_tokens( self, request: prediction_service.CountTokensRequest, @@ -235,12 +261,37 @@ async def post_count_tokens( ) -> prediction_service.CountTokensResponse: """Post-rpc interceptor for count_tokens - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_count_tokens_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_count_tokens` interceptor runs + before the `post_count_tokens_with_metadata` interceptor. """ return response + async def post_count_tokens_with_metadata( + self, + response: prediction_service.CountTokensResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.CountTokensResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for count_tokens + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_count_tokens_with_metadata` + interceptor in new development instead of the `post_count_tokens` interceptor. + When both interceptors are used, this `post_count_tokens_with_metadata` interceptor runs after the + `post_count_tokens` interceptor. The (possibly modified) response returned by + `post_count_tokens` will be passed to + `post_count_tokens_with_metadata`. + """ + return response, metadata + async def pre_direct_predict( self, request: prediction_service.DirectPredictRequest, @@ -260,12 +311,38 @@ async def post_direct_predict( ) -> prediction_service.DirectPredictResponse: """Post-rpc interceptor for direct_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_predict` interceptor runs + before the `post_direct_predict_with_metadata` interceptor. """ return response + async def post_direct_predict_with_metadata( + self, + response: prediction_service.DirectPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_predict_with_metadata` + interceptor in new development instead of the `post_direct_predict` interceptor. + When both interceptors are used, this `post_direct_predict_with_metadata` interceptor runs after the + `post_direct_predict` interceptor. The (possibly modified) response returned by + `post_direct_predict` will be passed to + `post_direct_predict_with_metadata`. + """ + return response, metadata + async def pre_direct_raw_predict( self, request: prediction_service.DirectRawPredictRequest, @@ -286,12 +363,38 @@ async def post_direct_raw_predict( ) -> prediction_service.DirectRawPredictResponse: """Post-rpc interceptor for direct_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_direct_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_direct_raw_predict` interceptor runs + before the `post_direct_raw_predict_with_metadata` interceptor. """ return response + async def post_direct_raw_predict_with_metadata( + self, + response: prediction_service.DirectRawPredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.DirectRawPredictResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for direct_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_direct_raw_predict_with_metadata` + interceptor in new development instead of the `post_direct_raw_predict` interceptor. + When both interceptors are used, this `post_direct_raw_predict_with_metadata` interceptor runs after the + `post_direct_raw_predict` interceptor. The (possibly modified) response returned by + `post_direct_raw_predict` will be passed to + `post_direct_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_explain( self, request: prediction_service.ExplainRequest, @@ -311,12 +414,37 @@ async def post_explain( ) -> prediction_service.ExplainResponse: """Post-rpc interceptor for explain - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_explain_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_explain` interceptor runs + before the `post_explain_with_metadata` interceptor. """ return response + async def post_explain_with_metadata( + self, + response: prediction_service.ExplainResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.ExplainResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for explain + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_explain_with_metadata` + interceptor in new development instead of the `post_explain` interceptor. + When both interceptors are used, this `post_explain_with_metadata` interceptor runs after the + `post_explain` interceptor. The (possibly modified) response returned by + `post_explain` will be passed to + `post_explain_with_metadata`. + """ + return response, metadata + async def pre_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -337,12 +465,38 @@ async def post_generate_content( ) -> prediction_service.GenerateContentResponse: """Post-rpc interceptor for generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_generate_content` interceptor runs + before the `post_generate_content_with_metadata` interceptor. """ return response + async def post_generate_content_with_metadata( + self, + response: prediction_service.GenerateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.GenerateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_generate_content_with_metadata` + interceptor in new development instead of the `post_generate_content` interceptor. + When both interceptors are used, this `post_generate_content_with_metadata` interceptor runs after the + `post_generate_content` interceptor. The (possibly modified) response returned by + `post_generate_content` will be passed to + `post_generate_content_with_metadata`. + """ + return response, metadata + async def pre_predict( self, request: prediction_service.PredictRequest, @@ -362,12 +516,37 @@ async def post_predict( ) -> prediction_service.PredictResponse: """Post-rpc interceptor for predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_predict` interceptor runs + before the `post_predict_with_metadata` interceptor. """ return response + async def post_predict_with_metadata( + self, + response: prediction_service.PredictResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + prediction_service.PredictResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_predict_with_metadata` + interceptor in new development instead of the `post_predict` interceptor. + When both interceptors are used, this `post_predict_with_metadata` interceptor runs after the + `post_predict` interceptor. The (possibly modified) response returned by + `post_predict` will be passed to + `post_predict_with_metadata`. + """ + return response, metadata + async def pre_raw_predict( self, request: prediction_service.RawPredictRequest, @@ -387,12 +566,35 @@ async def post_raw_predict( ) -> httpbody_pb2.HttpBody: """Post-rpc interceptor for raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_raw_predict` interceptor runs + before the `post_raw_predict_with_metadata` interceptor. """ return response + async def post_raw_predict_with_metadata( + self, + response: httpbody_pb2.HttpBody, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[httpbody_pb2.HttpBody, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_raw_predict_with_metadata` + interceptor in new development instead of the `post_raw_predict` interceptor. + When both interceptors are used, this `post_raw_predict_with_metadata` interceptor runs after the + `post_raw_predict` interceptor. The (possibly modified) response returned by + `post_raw_predict` will be passed to + `post_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_server_streaming_predict( self, request: prediction_service.StreamingPredictRequest, @@ -413,12 +615,38 @@ async def post_server_streaming_predict( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for server_streaming_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_server_streaming_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_server_streaming_predict` interceptor runs + before the `post_server_streaming_predict_with_metadata` interceptor. """ return response + async def post_server_streaming_predict_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for server_streaming_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_server_streaming_predict_with_metadata` + interceptor in new development instead of the `post_server_streaming_predict` interceptor. + When both interceptors are used, this `post_server_streaming_predict_with_metadata` interceptor runs after the + `post_server_streaming_predict` interceptor. The (possibly modified) response returned by + `post_server_streaming_predict` will be passed to + `post_server_streaming_predict_with_metadata`. + """ + return response, metadata + async def pre_stream_generate_content( self, request: prediction_service.GenerateContentRequest, @@ -439,12 +667,38 @@ async def post_stream_generate_content( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_generate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_generate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_generate_content` interceptor runs + before the `post_stream_generate_content_with_metadata` interceptor. """ return response + async def post_stream_generate_content_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_generate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_generate_content_with_metadata` + interceptor in new development instead of the `post_stream_generate_content` interceptor. + When both interceptors are used, this `post_stream_generate_content_with_metadata` interceptor runs after the + `post_stream_generate_content` interceptor. The (possibly modified) response returned by + `post_stream_generate_content` will be passed to + `post_stream_generate_content_with_metadata`. + """ + return response, metadata + async def pre_stream_raw_predict( self, request: prediction_service.StreamRawPredictRequest, @@ -465,12 +719,38 @@ async def post_stream_raw_predict( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_raw_predict - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_raw_predict_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the PredictionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_raw_predict` interceptor runs + before the `post_stream_raw_predict_with_metadata` interceptor. """ return response + async def post_stream_raw_predict_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_raw_predict + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the PredictionService server but before it is returned to user code. + + We recommend only using this `post_stream_raw_predict_with_metadata` + interceptor in new development instead of the `post_stream_raw_predict` interceptor. + When both interceptors are used, this `post_stream_raw_predict_with_metadata` interceptor runs after the + `post_stream_raw_predict` interceptor. The (possibly modified) response returned by + `post_stream_raw_predict` will be passed to + `post_stream_raw_predict_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1095,6 +1375,10 @@ async def __call__( response, httpbody_pb2.HttpBody ) resp = await self._interceptor.post_chat_completions(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_chat_completions_with_metadata( + resp, response_metadata + ) return resp class _CountTokens( @@ -1234,6 +1518,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_count_tokens(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_count_tokens_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1397,6 +1685,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_direct_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_direct_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1558,6 +1850,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_direct_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_direct_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1720,6 +2016,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_explain(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_explain_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1880,6 +2180,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_generate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2042,6 +2346,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2252,6 +2560,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_raw_predict_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2415,6 +2727,13 @@ async def __call__( response, prediction_service.StreamingPredictResponse ) resp = await self._interceptor.post_server_streaming_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_server_streaming_predict_with_metadata( + resp, response_metadata + ) return resp class _StreamDirectPredict( @@ -2588,6 +2907,13 @@ async def __call__( response, prediction_service.GenerateContentResponse ) resp = await self._interceptor.post_stream_generate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_stream_generate_content_with_metadata( + resp, response_metadata + ) return resp class _StreamingPredict( @@ -2809,6 +3135,10 @@ async def __call__( response, httpbody_pb2.HttpBody ) resp = await self._interceptor.post_stream_raw_predict(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stream_raw_predict_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/client.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/client.py index 04080fac7c..81f4cf1d8b 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -517,6 +519,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1041,16 +1070,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1096,16 +1129,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1268,16 +1305,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1389,16 +1430,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1511,16 +1556,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1571,16 +1620,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1626,16 +1679,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1681,16 +1738,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest.py index a29e58b162..ab06bcf899 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest.py @@ -121,12 +121,38 @@ def post_query_reasoning_engine( ) -> reasoning_engine_execution_service.QueryReasoningEngineResponse: """Post-rpc interceptor for query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_reasoning_engine` interceptor runs + before the `post_query_reasoning_engine_with_metadata` interceptor. """ return response + def post_query_reasoning_engine_with_metadata( + self, + response: reasoning_engine_execution_service.QueryReasoningEngineResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_execution_service.QueryReasoningEngineResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_query_reasoning_engine_with_metadata` interceptor runs after the + `post_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_query_reasoning_engine` will be passed to + `post_query_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_stream_query_reasoning_engine( self, request: reasoning_engine_execution_service.StreamQueryReasoningEngineRequest, @@ -147,12 +173,37 @@ def post_stream_query_reasoning_engine( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for stream_query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_query_reasoning_engine` interceptor runs + before the `post_stream_query_reasoning_engine_with_metadata` interceptor. """ return response + def post_stream_query_reasoning_engine_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for stream_query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_stream_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_stream_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_stream_query_reasoning_engine_with_metadata` interceptor runs after the + `post_stream_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_stream_query_reasoning_engine` will be passed to + `post_stream_query_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -621,6 +672,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_query_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -824,6 +879,13 @@ def __call__( resp = rest_streaming.ResponseIterator(response, httpbody_pb2.HttpBody) resp = self._interceptor.post_stream_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_stream_query_reasoning_engine_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest_asyncio.py index cdc52e80a7..d833f79ff6 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_execution_service/transports/rest_asyncio.py @@ -138,12 +138,38 @@ async def post_query_reasoning_engine( ) -> reasoning_engine_execution_service.QueryReasoningEngineResponse: """Post-rpc interceptor for query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_query_reasoning_engine` interceptor runs + before the `post_query_reasoning_engine_with_metadata` interceptor. """ return response + async def post_query_reasoning_engine_with_metadata( + self, + response: reasoning_engine_execution_service.QueryReasoningEngineResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_execution_service.QueryReasoningEngineResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_query_reasoning_engine_with_metadata` interceptor runs after the + `post_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_query_reasoning_engine` will be passed to + `post_query_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_stream_query_reasoning_engine( self, request: reasoning_engine_execution_service.StreamQueryReasoningEngineRequest, @@ -164,12 +190,38 @@ async def post_stream_query_reasoning_engine( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for stream_query_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stream_query_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineExecutionService server but before - it is returned to user code. + it is returned to user code. This `post_stream_query_reasoning_engine` interceptor runs + before the `post_stream_query_reasoning_engine_with_metadata` interceptor. """ return response + async def post_stream_query_reasoning_engine_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for stream_query_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineExecutionService server but before it is returned to user code. + + We recommend only using this `post_stream_query_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_stream_query_reasoning_engine` interceptor. + When both interceptors are used, this `post_stream_query_reasoning_engine_with_metadata` interceptor runs after the + `post_stream_query_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_stream_query_reasoning_engine` will be passed to + `post_stream_query_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -694,6 +746,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_query_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -907,6 +963,13 @@ async def __call__( response, httpbody_pb2.HttpBody ) resp = await self._interceptor.post_stream_query_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_stream_query_reasoning_engine_with_metadata( + resp, response_metadata + ) return resp @property diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/client.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/client.py index 5843e57107..3523435c57 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -513,6 +515,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1439,16 +1468,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1494,16 +1527,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1666,16 +1703,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1787,16 +1828,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1909,16 +1954,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1969,16 +2018,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2024,16 +2077,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2079,16 +2136,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest.py index 10aded0c2a..0a232956bd 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest.py @@ -146,12 +146,35 @@ def post_create_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_create_reasoning_engine` interceptor runs + before the `post_create_reasoning_engine_with_metadata` interceptor. """ return response + def post_create_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_create_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_create_reasoning_engine` interceptor. + When both interceptors are used, this `post_create_reasoning_engine_with_metadata` interceptor runs after the + `post_create_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_create_reasoning_engine` will be passed to + `post_create_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_delete_reasoning_engine( self, request: reasoning_engine_service.DeleteReasoningEngineRequest, @@ -172,12 +195,35 @@ def post_delete_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_reasoning_engine` interceptor runs + before the `post_delete_reasoning_engine_with_metadata` interceptor. """ return response + def post_delete_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_delete_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_delete_reasoning_engine` interceptor. + When both interceptors are used, this `post_delete_reasoning_engine_with_metadata` interceptor runs after the + `post_delete_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_delete_reasoning_engine` will be passed to + `post_delete_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_reasoning_engine( self, request: reasoning_engine_service.GetReasoningEngineRequest, @@ -198,12 +244,37 @@ def post_get_reasoning_engine( ) -> reasoning_engine.ReasoningEngine: """Post-rpc interceptor for get_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_get_reasoning_engine` interceptor runs + before the `post_get_reasoning_engine_with_metadata` interceptor. """ return response + def post_get_reasoning_engine_with_metadata( + self, + response: reasoning_engine.ReasoningEngine, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine.ReasoningEngine, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_get_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_get_reasoning_engine` interceptor. + When both interceptors are used, this `post_get_reasoning_engine_with_metadata` interceptor runs after the + `post_get_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_get_reasoning_engine` will be passed to + `post_get_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_list_reasoning_engines( self, request: reasoning_engine_service.ListReasoningEnginesRequest, @@ -224,12 +295,38 @@ def post_list_reasoning_engines( ) -> reasoning_engine_service.ListReasoningEnginesResponse: """Post-rpc interceptor for list_reasoning_engines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_reasoning_engines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_list_reasoning_engines` interceptor runs + before the `post_list_reasoning_engines_with_metadata` interceptor. """ return response + def post_list_reasoning_engines_with_metadata( + self, + response: reasoning_engine_service.ListReasoningEnginesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_service.ListReasoningEnginesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_reasoning_engines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_list_reasoning_engines_with_metadata` + interceptor in new development instead of the `post_list_reasoning_engines` interceptor. + When both interceptors are used, this `post_list_reasoning_engines_with_metadata` interceptor runs after the + `post_list_reasoning_engines` interceptor. The (possibly modified) response returned by + `post_list_reasoning_engines` will be passed to + `post_list_reasoning_engines_with_metadata`. + """ + return response, metadata + def pre_update_reasoning_engine( self, request: reasoning_engine_service.UpdateReasoningEngineRequest, @@ -250,12 +347,35 @@ def post_update_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_update_reasoning_engine` interceptor runs + before the `post_update_reasoning_engine_with_metadata` interceptor. """ return response + def post_update_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_update_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_update_reasoning_engine` interceptor. + When both interceptors are used, this `post_update_reasoning_engine_with_metadata` interceptor runs after the + `post_update_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_update_reasoning_engine` will be passed to + `post_update_reasoning_engine_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2802,6 +2922,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2948,6 +3072,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3099,6 +3227,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3250,6 +3382,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_reasoning_engines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_reasoning_engines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3406,6 +3542,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest_asyncio.py index 56611a471e..f8cf3c4a0a 100644 --- a/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/reasoning_engine_service/transports/rest_asyncio.py @@ -164,12 +164,35 @@ async def post_create_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_create_reasoning_engine` interceptor runs + before the `post_create_reasoning_engine_with_metadata` interceptor. """ return response + async def post_create_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_create_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_create_reasoning_engine` interceptor. + When both interceptors are used, this `post_create_reasoning_engine_with_metadata` interceptor runs after the + `post_create_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_create_reasoning_engine` will be passed to + `post_create_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_delete_reasoning_engine( self, request: reasoning_engine_service.DeleteReasoningEngineRequest, @@ -190,12 +213,35 @@ async def post_delete_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_delete_reasoning_engine` interceptor runs + before the `post_delete_reasoning_engine_with_metadata` interceptor. """ return response + async def post_delete_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_delete_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_delete_reasoning_engine` interceptor. + When both interceptors are used, this `post_delete_reasoning_engine_with_metadata` interceptor runs after the + `post_delete_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_delete_reasoning_engine` will be passed to + `post_delete_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_reasoning_engine( self, request: reasoning_engine_service.GetReasoningEngineRequest, @@ -216,12 +262,37 @@ async def post_get_reasoning_engine( ) -> reasoning_engine.ReasoningEngine: """Post-rpc interceptor for get_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_get_reasoning_engine` interceptor runs + before the `post_get_reasoning_engine_with_metadata` interceptor. """ return response + async def post_get_reasoning_engine_with_metadata( + self, + response: reasoning_engine.ReasoningEngine, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine.ReasoningEngine, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for get_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_get_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_get_reasoning_engine` interceptor. + When both interceptors are used, this `post_get_reasoning_engine_with_metadata` interceptor runs after the + `post_get_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_get_reasoning_engine` will be passed to + `post_get_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_list_reasoning_engines( self, request: reasoning_engine_service.ListReasoningEnginesRequest, @@ -242,12 +313,38 @@ async def post_list_reasoning_engines( ) -> reasoning_engine_service.ListReasoningEnginesResponse: """Post-rpc interceptor for list_reasoning_engines - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_reasoning_engines_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_list_reasoning_engines` interceptor runs + before the `post_list_reasoning_engines_with_metadata` interceptor. """ return response + async def post_list_reasoning_engines_with_metadata( + self, + response: reasoning_engine_service.ListReasoningEnginesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + reasoning_engine_service.ListReasoningEnginesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_reasoning_engines + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_list_reasoning_engines_with_metadata` + interceptor in new development instead of the `post_list_reasoning_engines` interceptor. + When both interceptors are used, this `post_list_reasoning_engines_with_metadata` interceptor runs after the + `post_list_reasoning_engines` interceptor. The (possibly modified) response returned by + `post_list_reasoning_engines` will be passed to + `post_list_reasoning_engines_with_metadata`. + """ + return response, metadata + async def pre_update_reasoning_engine( self, request: reasoning_engine_service.UpdateReasoningEngineRequest, @@ -268,12 +365,35 @@ async def post_update_reasoning_engine( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_reasoning_engine - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_reasoning_engine_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ReasoningEngineService server but before - it is returned to user code. + it is returned to user code. This `post_update_reasoning_engine` interceptor runs + before the `post_update_reasoning_engine_with_metadata` interceptor. """ return response + async def post_update_reasoning_engine_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_reasoning_engine + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ReasoningEngineService server but before it is returned to user code. + + We recommend only using this `post_update_reasoning_engine_with_metadata` + interceptor in new development instead of the `post_update_reasoning_engine` interceptor. + When both interceptors are used, this `post_update_reasoning_engine_with_metadata` interceptor runs after the + `post_update_reasoning_engine` interceptor. The (possibly modified) response returned by + `post_update_reasoning_engine` will be passed to + `post_update_reasoning_engine_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -811,6 +931,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -967,6 +1094,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1122,6 +1256,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1277,6 +1415,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_reasoning_engines(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_reasoning_engines_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1443,6 +1585,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_reasoning_engine(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_reasoning_engine_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/schedule_service/client.py b/google/cloud/aiplatform_v1beta1/services/schedule_service/client.py index 084a9b4c91..42b8ffcb35 100644 --- a/google/cloud/aiplatform_v1beta1/services/schedule_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/schedule_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -872,6 +874,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1981,16 +2010,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2036,16 +2069,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2208,16 +2245,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2329,16 +2370,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2451,16 +2496,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2511,16 +2560,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2566,16 +2619,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2621,16 +2678,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest.py index 43445b1a1c..76d22ec5a2 100644 --- a/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest.py @@ -155,12 +155,35 @@ def post_create_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for create_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_create_schedule` interceptor runs + before the `post_create_schedule_with_metadata` interceptor. """ return response + def post_create_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_create_schedule_with_metadata` + interceptor in new development instead of the `post_create_schedule` interceptor. + When both interceptors are used, this `post_create_schedule_with_metadata` interceptor runs after the + `post_create_schedule` interceptor. The (possibly modified) response returned by + `post_create_schedule` will be passed to + `post_create_schedule_with_metadata`. + """ + return response, metadata + def pre_delete_schedule( self, request: schedule_service.DeleteScheduleRequest, @@ -180,12 +203,35 @@ def post_delete_schedule( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_delete_schedule` interceptor runs + before the `post_delete_schedule_with_metadata` interceptor. """ return response + def post_delete_schedule_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_delete_schedule_with_metadata` + interceptor in new development instead of the `post_delete_schedule` interceptor. + When both interceptors are used, this `post_delete_schedule_with_metadata` interceptor runs after the + `post_delete_schedule` interceptor. The (possibly modified) response returned by + `post_delete_schedule` will be passed to + `post_delete_schedule_with_metadata`. + """ + return response, metadata + def pre_get_schedule( self, request: schedule_service.GetScheduleRequest, @@ -203,12 +249,35 @@ def pre_get_schedule( def post_get_schedule(self, response: schedule.Schedule) -> schedule.Schedule: """Post-rpc interceptor for get_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_get_schedule` interceptor runs + before the `post_get_schedule_with_metadata` interceptor. """ return response + def post_get_schedule_with_metadata( + self, + response: schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_get_schedule_with_metadata` + interceptor in new development instead of the `post_get_schedule` interceptor. + When both interceptors are used, this `post_get_schedule_with_metadata` interceptor runs after the + `post_get_schedule` interceptor. The (possibly modified) response returned by + `post_get_schedule` will be passed to + `post_get_schedule_with_metadata`. + """ + return response, metadata + def pre_list_schedules( self, request: schedule_service.ListSchedulesRequest, @@ -228,12 +297,37 @@ def post_list_schedules( ) -> schedule_service.ListSchedulesResponse: """Post-rpc interceptor for list_schedules - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_schedules_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_list_schedules` interceptor runs + before the `post_list_schedules_with_metadata` interceptor. """ return response + def post_list_schedules_with_metadata( + self, + response: schedule_service.ListSchedulesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + schedule_service.ListSchedulesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_schedules + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_list_schedules_with_metadata` + interceptor in new development instead of the `post_list_schedules` interceptor. + When both interceptors are used, this `post_list_schedules_with_metadata` interceptor runs after the + `post_list_schedules` interceptor. The (possibly modified) response returned by + `post_list_schedules` will be passed to + `post_list_schedules_with_metadata`. + """ + return response, metadata + def pre_pause_schedule( self, request: schedule_service.PauseScheduleRequest, @@ -281,12 +375,35 @@ def post_update_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for update_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_update_schedule` interceptor runs + before the `post_update_schedule_with_metadata` interceptor. """ return response + def post_update_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_update_schedule_with_metadata` + interceptor in new development instead of the `post_update_schedule` interceptor. + When both interceptors are used, this `post_update_schedule_with_metadata` interceptor runs after the + `post_update_schedule` interceptor. The (possibly modified) response returned by + `post_update_schedule` will be passed to + `post_update_schedule_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2835,6 +2952,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2978,6 +3099,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3124,6 +3249,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3268,6 +3397,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_schedules(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_schedules_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3650,6 +3783,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest_asyncio.py index b9dbfde87a..668e10d3cf 100644 --- a/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/schedule_service/transports/rest_asyncio.py @@ -173,12 +173,35 @@ async def post_create_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for create_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_create_schedule` interceptor runs + before the `post_create_schedule_with_metadata` interceptor. """ return response + async def post_create_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_create_schedule_with_metadata` + interceptor in new development instead of the `post_create_schedule` interceptor. + When both interceptors are used, this `post_create_schedule_with_metadata` interceptor runs after the + `post_create_schedule` interceptor. The (possibly modified) response returned by + `post_create_schedule` will be passed to + `post_create_schedule_with_metadata`. + """ + return response, metadata + async def pre_delete_schedule( self, request: schedule_service.DeleteScheduleRequest, @@ -198,12 +221,35 @@ async def post_delete_schedule( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_delete_schedule` interceptor runs + before the `post_delete_schedule_with_metadata` interceptor. """ return response + async def post_delete_schedule_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_delete_schedule_with_metadata` + interceptor in new development instead of the `post_delete_schedule` interceptor. + When both interceptors are used, this `post_delete_schedule_with_metadata` interceptor runs after the + `post_delete_schedule` interceptor. The (possibly modified) response returned by + `post_delete_schedule` will be passed to + `post_delete_schedule_with_metadata`. + """ + return response, metadata + async def pre_get_schedule( self, request: schedule_service.GetScheduleRequest, @@ -221,12 +267,35 @@ async def pre_get_schedule( async def post_get_schedule(self, response: schedule.Schedule) -> schedule.Schedule: """Post-rpc interceptor for get_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_get_schedule` interceptor runs + before the `post_get_schedule_with_metadata` interceptor. """ return response + async def post_get_schedule_with_metadata( + self, + response: schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_get_schedule_with_metadata` + interceptor in new development instead of the `post_get_schedule` interceptor. + When both interceptors are used, this `post_get_schedule_with_metadata` interceptor runs after the + `post_get_schedule` interceptor. The (possibly modified) response returned by + `post_get_schedule` will be passed to + `post_get_schedule_with_metadata`. + """ + return response, metadata + async def pre_list_schedules( self, request: schedule_service.ListSchedulesRequest, @@ -246,12 +315,37 @@ async def post_list_schedules( ) -> schedule_service.ListSchedulesResponse: """Post-rpc interceptor for list_schedules - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_schedules_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_list_schedules` interceptor runs + before the `post_list_schedules_with_metadata` interceptor. """ return response + async def post_list_schedules_with_metadata( + self, + response: schedule_service.ListSchedulesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + schedule_service.ListSchedulesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_schedules + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_list_schedules_with_metadata` + interceptor in new development instead of the `post_list_schedules` interceptor. + When both interceptors are used, this `post_list_schedules_with_metadata` interceptor runs after the + `post_list_schedules` interceptor. The (possibly modified) response returned by + `post_list_schedules` will be passed to + `post_list_schedules_with_metadata`. + """ + return response, metadata + async def pre_pause_schedule( self, request: schedule_service.PauseScheduleRequest, @@ -299,12 +393,35 @@ async def post_update_schedule( ) -> gca_schedule.Schedule: """Post-rpc interceptor for update_schedule - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_schedule_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the ScheduleService server but before - it is returned to user code. + it is returned to user code. This `post_update_schedule` interceptor runs + before the `post_update_schedule_with_metadata` interceptor. """ return response + async def post_update_schedule_with_metadata( + self, + response: gca_schedule.Schedule, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_schedule.Schedule, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_schedule + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the ScheduleService server but before it is returned to user code. + + We recommend only using this `post_update_schedule_with_metadata` + interceptor in new development instead of the `post_update_schedule` interceptor. + When both interceptors are used, this `post_update_schedule_with_metadata` interceptor runs after the + `post_update_schedule` interceptor. The (possibly modified) response returned by + `post_update_schedule` will be passed to + `post_update_schedule_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -853,6 +970,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1009,6 +1130,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1165,6 +1290,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1320,6 +1449,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_schedules(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_schedules_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1735,6 +1868,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_schedule(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_schedule_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/client.py b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/client.py index c448e2dd57..0811e6b052 100644 --- a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -516,6 +518,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1457,16 +1486,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1512,16 +1545,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1684,16 +1721,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1805,16 +1846,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1927,16 +1972,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1987,16 +2036,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2042,16 +2095,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2097,16 +2154,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest.py index c6ce7016fe..60bfa8b937 100644 --- a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest.py @@ -146,12 +146,35 @@ def post_create_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_specialist_pool` interceptor runs + before the `post_create_specialist_pool_with_metadata` interceptor. """ return response + def post_create_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_create_specialist_pool_with_metadata` + interceptor in new development instead of the `post_create_specialist_pool` interceptor. + When both interceptors are used, this `post_create_specialist_pool_with_metadata` interceptor runs after the + `post_create_specialist_pool` interceptor. The (possibly modified) response returned by + `post_create_specialist_pool` will be passed to + `post_create_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_delete_specialist_pool( self, request: specialist_pool_service.DeleteSpecialistPoolRequest, @@ -172,12 +195,35 @@ def post_delete_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_specialist_pool` interceptor runs + before the `post_delete_specialist_pool_with_metadata` interceptor. """ return response + def post_delete_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_delete_specialist_pool_with_metadata` + interceptor in new development instead of the `post_delete_specialist_pool` interceptor. + When both interceptors are used, this `post_delete_specialist_pool_with_metadata` interceptor runs after the + `post_delete_specialist_pool` interceptor. The (possibly modified) response returned by + `post_delete_specialist_pool` will be passed to + `post_delete_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_get_specialist_pool( self, request: specialist_pool_service.GetSpecialistPoolRequest, @@ -198,12 +244,35 @@ def post_get_specialist_pool( ) -> specialist_pool.SpecialistPool: """Post-rpc interceptor for get_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_specialist_pool` interceptor runs + before the `post_get_specialist_pool_with_metadata` interceptor. """ return response + def post_get_specialist_pool_with_metadata( + self, + response: specialist_pool.SpecialistPool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[specialist_pool.SpecialistPool, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_get_specialist_pool_with_metadata` + interceptor in new development instead of the `post_get_specialist_pool` interceptor. + When both interceptors are used, this `post_get_specialist_pool_with_metadata` interceptor runs after the + `post_get_specialist_pool` interceptor. The (possibly modified) response returned by + `post_get_specialist_pool` will be passed to + `post_get_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_list_specialist_pools( self, request: specialist_pool_service.ListSpecialistPoolsRequest, @@ -224,12 +293,38 @@ def post_list_specialist_pools( ) -> specialist_pool_service.ListSpecialistPoolsResponse: """Post-rpc interceptor for list_specialist_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_specialist_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_specialist_pools` interceptor runs + before the `post_list_specialist_pools_with_metadata` interceptor. """ return response + def post_list_specialist_pools_with_metadata( + self, + response: specialist_pool_service.ListSpecialistPoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + specialist_pool_service.ListSpecialistPoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_specialist_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_list_specialist_pools_with_metadata` + interceptor in new development instead of the `post_list_specialist_pools` interceptor. + When both interceptors are used, this `post_list_specialist_pools_with_metadata` interceptor runs after the + `post_list_specialist_pools` interceptor. The (possibly modified) response returned by + `post_list_specialist_pools` will be passed to + `post_list_specialist_pools_with_metadata`. + """ + return response, metadata + def pre_update_specialist_pool( self, request: specialist_pool_service.UpdateSpecialistPoolRequest, @@ -250,12 +345,35 @@ def post_update_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_specialist_pool` interceptor runs + before the `post_update_specialist_pool_with_metadata` interceptor. """ return response + def post_update_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_update_specialist_pool_with_metadata` + interceptor in new development instead of the `post_update_specialist_pool` interceptor. + When both interceptors are used, this `post_update_specialist_pool_with_metadata` interceptor runs after the + `post_update_specialist_pool` interceptor. The (possibly modified) response returned by + `post_update_specialist_pool` will be passed to + `post_update_specialist_pool_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2809,6 +2927,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2957,6 +3079,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3116,6 +3242,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3265,6 +3395,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_specialist_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_specialist_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3423,6 +3557,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest_asyncio.py index 19d838e723..ac9f45fc61 100644 --- a/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/specialist_pool_service/transports/rest_asyncio.py @@ -164,12 +164,35 @@ async def post_create_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_create_specialist_pool` interceptor runs + before the `post_create_specialist_pool_with_metadata` interceptor. """ return response + async def post_create_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_create_specialist_pool_with_metadata` + interceptor in new development instead of the `post_create_specialist_pool` interceptor. + When both interceptors are used, this `post_create_specialist_pool_with_metadata` interceptor runs after the + `post_create_specialist_pool` interceptor. The (possibly modified) response returned by + `post_create_specialist_pool` will be passed to + `post_create_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_delete_specialist_pool( self, request: specialist_pool_service.DeleteSpecialistPoolRequest, @@ -190,12 +213,35 @@ async def post_delete_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_delete_specialist_pool` interceptor runs + before the `post_delete_specialist_pool_with_metadata` interceptor. """ return response + async def post_delete_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_delete_specialist_pool_with_metadata` + interceptor in new development instead of the `post_delete_specialist_pool` interceptor. + When both interceptors are used, this `post_delete_specialist_pool_with_metadata` interceptor runs after the + `post_delete_specialist_pool` interceptor. The (possibly modified) response returned by + `post_delete_specialist_pool` will be passed to + `post_delete_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_get_specialist_pool( self, request: specialist_pool_service.GetSpecialistPoolRequest, @@ -216,12 +262,35 @@ async def post_get_specialist_pool( ) -> specialist_pool.SpecialistPool: """Post-rpc interceptor for get_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_get_specialist_pool` interceptor runs + before the `post_get_specialist_pool_with_metadata` interceptor. """ return response + async def post_get_specialist_pool_with_metadata( + self, + response: specialist_pool.SpecialistPool, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[specialist_pool.SpecialistPool, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_get_specialist_pool_with_metadata` + interceptor in new development instead of the `post_get_specialist_pool` interceptor. + When both interceptors are used, this `post_get_specialist_pool_with_metadata` interceptor runs after the + `post_get_specialist_pool` interceptor. The (possibly modified) response returned by + `post_get_specialist_pool` will be passed to + `post_get_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_list_specialist_pools( self, request: specialist_pool_service.ListSpecialistPoolsRequest, @@ -242,12 +311,38 @@ async def post_list_specialist_pools( ) -> specialist_pool_service.ListSpecialistPoolsResponse: """Post-rpc interceptor for list_specialist_pools - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_specialist_pools_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_list_specialist_pools` interceptor runs + before the `post_list_specialist_pools_with_metadata` interceptor. """ return response + async def post_list_specialist_pools_with_metadata( + self, + response: specialist_pool_service.ListSpecialistPoolsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + specialist_pool_service.ListSpecialistPoolsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_specialist_pools + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_list_specialist_pools_with_metadata` + interceptor in new development instead of the `post_list_specialist_pools` interceptor. + When both interceptors are used, this `post_list_specialist_pools_with_metadata` interceptor runs after the + `post_list_specialist_pools` interceptor. The (possibly modified) response returned by + `post_list_specialist_pools` will be passed to + `post_list_specialist_pools_with_metadata`. + """ + return response, metadata + async def pre_update_specialist_pool( self, request: specialist_pool_service.UpdateSpecialistPoolRequest, @@ -268,12 +363,35 @@ async def post_update_specialist_pool( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_specialist_pool - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_specialist_pool_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the SpecialistPoolService server but before - it is returned to user code. + it is returned to user code. This `post_update_specialist_pool` interceptor runs + before the `post_update_specialist_pool_with_metadata` interceptor. """ return response + async def post_update_specialist_pool_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_specialist_pool + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the SpecialistPoolService server but before it is returned to user code. + + We recommend only using this `post_update_specialist_pool_with_metadata` + interceptor in new development instead of the `post_update_specialist_pool` interceptor. + When both interceptors are used, this `post_update_specialist_pool_with_metadata` interceptor runs after the + `post_update_specialist_pool` interceptor. The (possibly modified) response returned by + `post_update_specialist_pool` will be passed to + `post_update_specialist_pool_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -812,6 +930,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -966,6 +1088,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1129,6 +1255,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1282,6 +1412,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_specialist_pools(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_specialist_pools_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1446,6 +1580,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_specialist_pool(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_specialist_pool_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/client.py b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/client.py index d23e6fe17d..76a2bf785a 100644 --- a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -602,6 +604,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -4706,16 +4735,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -4761,16 +4794,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -4933,16 +4970,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -5054,16 +5095,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -5176,16 +5221,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -5236,16 +5285,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -5291,16 +5344,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -5346,16 +5403,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest.py index fe6cba2856..5d74e52764 100644 --- a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest.py @@ -356,12 +356,38 @@ def post_batch_create_tensorboard_runs( ) -> tensorboard_service.BatchCreateTensorboardRunsResponse: """Post-rpc interceptor for batch_create_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_runs` interceptor runs + before the `post_batch_create_tensorboard_runs_with_metadata` interceptor. """ return response + def post_batch_create_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_runs` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_runs_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_runs` will be passed to + `post_batch_create_tensorboard_runs_with_metadata`. + """ + return response, metadata + def pre_batch_create_tensorboard_time_series( self, request: tensorboard_service.BatchCreateTensorboardTimeSeriesRequest, @@ -382,12 +408,38 @@ def post_batch_create_tensorboard_time_series( ) -> tensorboard_service.BatchCreateTensorboardTimeSeriesResponse: """Post-rpc interceptor for batch_create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_time_series` interceptor runs + before the `post_batch_create_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_batch_create_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_time_series` will be passed to + `post_batch_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_batch_read_tensorboard_time_series_data( self, request: tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest, @@ -408,12 +460,38 @@ def post_batch_read_tensorboard_time_series_data( ) -> tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for batch_read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_tensorboard_time_series_data` interceptor runs + before the `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_batch_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_batch_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_batch_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_batch_read_tensorboard_time_series_data` will be passed to + `post_batch_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard( self, request: tensorboard_service.CreateTensorboardRequest, @@ -434,12 +512,35 @@ def post_create_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard` interceptor runs + before the `post_create_tensorboard_with_metadata` interceptor. """ return response + def post_create_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_with_metadata` + interceptor in new development instead of the `post_create_tensorboard` interceptor. + When both interceptors are used, this `post_create_tensorboard_with_metadata` interceptor runs after the + `post_create_tensorboard` interceptor. The (possibly modified) response returned by + `post_create_tensorboard` will be passed to + `post_create_tensorboard_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_experiment( self, request: tensorboard_service.CreateTensorboardExperimentRequest, @@ -460,12 +561,38 @@ def post_create_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for create_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_experiment` interceptor runs + before the `post_create_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_create_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_create_tensorboard_experiment_with_metadata` interceptor runs after the + `post_create_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_experiment` will be passed to + `post_create_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_run( self, request: tensorboard_service.CreateTensorboardRunRequest, @@ -486,12 +613,37 @@ def post_create_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for create_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_run` interceptor runs + before the `post_create_tensorboard_run_with_metadata` interceptor. """ return response + def post_create_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_run` interceptor. + When both interceptors are used, this `post_create_tensorboard_run_with_metadata` interceptor runs after the + `post_create_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_run` will be passed to + `post_create_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_create_tensorboard_time_series( self, request: tensorboard_service.CreateTensorboardTimeSeriesRequest, @@ -512,12 +664,38 @@ def post_create_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_time_series` interceptor runs + before the `post_create_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_create_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_time_series` will be passed to + `post_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard( self, request: tensorboard_service.DeleteTensorboardRequest, @@ -538,12 +716,35 @@ def post_delete_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard` interceptor runs + before the `post_delete_tensorboard_with_metadata` interceptor. """ return response + def post_delete_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard` interceptor. + When both interceptors are used, this `post_delete_tensorboard_with_metadata` interceptor runs after the + `post_delete_tensorboard` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard` will be passed to + `post_delete_tensorboard_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_experiment( self, request: tensorboard_service.DeleteTensorboardExperimentRequest, @@ -564,12 +765,35 @@ def post_delete_tensorboard_experiment( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_experiment` interceptor runs + before the `post_delete_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_delete_tensorboard_experiment_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_delete_tensorboard_experiment_with_metadata` interceptor runs after the + `post_delete_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_experiment` will be passed to + `post_delete_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_run( self, request: tensorboard_service.DeleteTensorboardRunRequest, @@ -590,12 +814,35 @@ def post_delete_tensorboard_run( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_run` interceptor runs + before the `post_delete_tensorboard_run_with_metadata` interceptor. """ return response + def post_delete_tensorboard_run_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_run` interceptor. + When both interceptors are used, this `post_delete_tensorboard_run_with_metadata` interceptor runs after the + `post_delete_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_run` will be passed to + `post_delete_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_delete_tensorboard_time_series( self, request: tensorboard_service.DeleteTensorboardTimeSeriesRequest, @@ -616,12 +863,35 @@ def post_delete_tensorboard_time_series( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_time_series` interceptor runs + before the `post_delete_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_delete_tensorboard_time_series_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_delete_tensorboard_time_series_with_metadata` interceptor runs after the + `post_delete_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_time_series` will be passed to + `post_delete_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_export_tensorboard_time_series_data( self, request: tensorboard_service.ExportTensorboardTimeSeriesDataRequest, @@ -642,12 +912,38 @@ def post_export_tensorboard_time_series_data( ) -> tensorboard_service.ExportTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for export_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_export_tensorboard_time_series_data` interceptor runs + before the `post_export_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_export_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for export_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_export_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_export_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_export_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_export_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_export_tensorboard_time_series_data` will be passed to + `post_export_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard( self, request: tensorboard_service.GetTensorboardRequest, @@ -668,12 +964,35 @@ def post_get_tensorboard( ) -> tensorboard.Tensorboard: """Post-rpc interceptor for get_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard` interceptor runs + before the `post_get_tensorboard_with_metadata` interceptor. """ return response + def post_get_tensorboard_with_metadata( + self, + response: tensorboard.Tensorboard, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard.Tensorboard, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_with_metadata` + interceptor in new development instead of the `post_get_tensorboard` interceptor. + When both interceptors are used, this `post_get_tensorboard_with_metadata` interceptor runs after the + `post_get_tensorboard` interceptor. The (possibly modified) response returned by + `post_get_tensorboard` will be passed to + `post_get_tensorboard_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_experiment( self, request: tensorboard_service.GetTensorboardExperimentRequest, @@ -694,12 +1013,38 @@ def post_get_tensorboard_experiment( ) -> tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for get_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_experiment` interceptor runs + before the `post_get_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_get_tensorboard_experiment_with_metadata( + self, + response: tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_get_tensorboard_experiment_with_metadata` interceptor runs after the + `post_get_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_experiment` will be passed to + `post_get_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_run( self, request: tensorboard_service.GetTensorboardRunRequest, @@ -720,12 +1065,35 @@ def post_get_tensorboard_run( ) -> tensorboard_run.TensorboardRun: """Post-rpc interceptor for get_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_run` interceptor runs + before the `post_get_tensorboard_run_with_metadata` interceptor. """ return response + def post_get_tensorboard_run_with_metadata( + self, + response: tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_run` interceptor. + When both interceptors are used, this `post_get_tensorboard_run_with_metadata` interceptor runs after the + `post_get_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_run` will be passed to + `post_get_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_get_tensorboard_time_series( self, request: tensorboard_service.GetTensorboardTimeSeriesRequest, @@ -746,12 +1114,38 @@ def post_get_tensorboard_time_series( ) -> tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for get_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_time_series` interceptor runs + before the `post_get_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_get_tensorboard_time_series_with_metadata( + self, + response: tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_get_tensorboard_time_series_with_metadata` interceptor runs after the + `post_get_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_time_series` will be passed to + `post_get_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_experiments( self, request: tensorboard_service.ListTensorboardExperimentsRequest, @@ -772,12 +1166,38 @@ def post_list_tensorboard_experiments( ) -> tensorboard_service.ListTensorboardExperimentsResponse: """Post-rpc interceptor for list_tensorboard_experiments - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_experiments_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_experiments` interceptor runs + before the `post_list_tensorboard_experiments_with_metadata` interceptor. """ return response + def post_list_tensorboard_experiments_with_metadata( + self, + response: tensorboard_service.ListTensorboardExperimentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardExperimentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_experiments + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_experiments_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_experiments` interceptor. + When both interceptors are used, this `post_list_tensorboard_experiments_with_metadata` interceptor runs after the + `post_list_tensorboard_experiments` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_experiments` will be passed to + `post_list_tensorboard_experiments_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_runs( self, request: tensorboard_service.ListTensorboardRunsRequest, @@ -798,12 +1218,38 @@ def post_list_tensorboard_runs( ) -> tensorboard_service.ListTensorboardRunsResponse: """Post-rpc interceptor for list_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_runs` interceptor runs + before the `post_list_tensorboard_runs_with_metadata` interceptor. """ return response + def post_list_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.ListTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_runs` interceptor. + When both interceptors are used, this `post_list_tensorboard_runs_with_metadata` interceptor runs after the + `post_list_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_runs` will be passed to + `post_list_tensorboard_runs_with_metadata`. + """ + return response, metadata + def pre_list_tensorboards( self, request: tensorboard_service.ListTensorboardsRequest, @@ -824,12 +1270,38 @@ def post_list_tensorboards( ) -> tensorboard_service.ListTensorboardsResponse: """Post-rpc interceptor for list_tensorboards - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboards_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboards` interceptor runs + before the `post_list_tensorboards_with_metadata` interceptor. """ return response + def post_list_tensorboards_with_metadata( + self, + response: tensorboard_service.ListTensorboardsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboards + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboards_with_metadata` + interceptor in new development instead of the `post_list_tensorboards` interceptor. + When both interceptors are used, this `post_list_tensorboards_with_metadata` interceptor runs after the + `post_list_tensorboards` interceptor. The (possibly modified) response returned by + `post_list_tensorboards` will be passed to + `post_list_tensorboards_with_metadata`. + """ + return response, metadata + def pre_list_tensorboard_time_series( self, request: tensorboard_service.ListTensorboardTimeSeriesRequest, @@ -850,12 +1322,38 @@ def post_list_tensorboard_time_series( ) -> tensorboard_service.ListTensorboardTimeSeriesResponse: """Post-rpc interceptor for list_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_time_series` interceptor runs + before the `post_list_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_list_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.ListTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_list_tensorboard_time_series_with_metadata` interceptor runs after the + `post_list_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_time_series` will be passed to + `post_list_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_blob_data( self, request: tensorboard_service.ReadTensorboardBlobDataRequest, @@ -876,12 +1374,37 @@ def post_read_tensorboard_blob_data( ) -> rest_streaming.ResponseIterator: """Post-rpc interceptor for read_tensorboard_blob_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_blob_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_blob_data` interceptor runs + before the `post_read_tensorboard_blob_data_with_metadata` interceptor. """ return response + def post_read_tensorboard_blob_data_with_metadata( + self, + response: rest_streaming.ResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming.ResponseIterator, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for read_tensorboard_blob_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_blob_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_blob_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_blob_data_with_metadata` interceptor runs after the + `post_read_tensorboard_blob_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_blob_data` will be passed to + `post_read_tensorboard_blob_data_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_size( self, request: tensorboard_service.ReadTensorboardSizeRequest, @@ -902,12 +1425,38 @@ def post_read_tensorboard_size( ) -> tensorboard_service.ReadTensorboardSizeResponse: """Post-rpc interceptor for read_tensorboard_size - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_size_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_size` interceptor runs + before the `post_read_tensorboard_size_with_metadata` interceptor. """ return response + def post_read_tensorboard_size_with_metadata( + self, + response: tensorboard_service.ReadTensorboardSizeResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardSizeResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_size + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_size_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_size` interceptor. + When both interceptors are used, this `post_read_tensorboard_size_with_metadata` interceptor runs after the + `post_read_tensorboard_size` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_size` will be passed to + `post_read_tensorboard_size_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_time_series_data( self, request: tensorboard_service.ReadTensorboardTimeSeriesDataRequest, @@ -928,12 +1477,38 @@ def post_read_tensorboard_time_series_data( ) -> tensorboard_service.ReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_time_series_data` interceptor runs + before the `post_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + def post_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_time_series_data` will be passed to + `post_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + def pre_read_tensorboard_usage( self, request: tensorboard_service.ReadTensorboardUsageRequest, @@ -954,12 +1529,38 @@ def post_read_tensorboard_usage( ) -> tensorboard_service.ReadTensorboardUsageResponse: """Post-rpc interceptor for read_tensorboard_usage - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_usage_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_usage` interceptor runs + before the `post_read_tensorboard_usage_with_metadata` interceptor. """ return response + def post_read_tensorboard_usage_with_metadata( + self, + response: tensorboard_service.ReadTensorboardUsageResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardUsageResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_usage + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_usage_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_usage` interceptor. + When both interceptors are used, this `post_read_tensorboard_usage_with_metadata` interceptor runs after the + `post_read_tensorboard_usage` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_usage` will be passed to + `post_read_tensorboard_usage_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard( self, request: tensorboard_service.UpdateTensorboardRequest, @@ -980,12 +1581,35 @@ def post_update_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard` interceptor runs + before the `post_update_tensorboard_with_metadata` interceptor. """ return response + def post_update_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_with_metadata` + interceptor in new development instead of the `post_update_tensorboard` interceptor. + When both interceptors are used, this `post_update_tensorboard_with_metadata` interceptor runs after the + `post_update_tensorboard` interceptor. The (possibly modified) response returned by + `post_update_tensorboard` will be passed to + `post_update_tensorboard_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_experiment( self, request: tensorboard_service.UpdateTensorboardExperimentRequest, @@ -1006,12 +1630,38 @@ def post_update_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for update_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_experiment` interceptor runs + before the `post_update_tensorboard_experiment_with_metadata` interceptor. """ return response + def post_update_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_update_tensorboard_experiment_with_metadata` interceptor runs after the + `post_update_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_experiment` will be passed to + `post_update_tensorboard_experiment_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_run( self, request: tensorboard_service.UpdateTensorboardRunRequest, @@ -1032,12 +1682,37 @@ def post_update_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for update_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_run` interceptor runs + before the `post_update_tensorboard_run_with_metadata` interceptor. """ return response + def post_update_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_run` interceptor. + When both interceptors are used, this `post_update_tensorboard_run_with_metadata` interceptor runs after the + `post_update_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_run` will be passed to + `post_update_tensorboard_run_with_metadata`. + """ + return response, metadata + def pre_update_tensorboard_time_series( self, request: tensorboard_service.UpdateTensorboardTimeSeriesRequest, @@ -1058,12 +1733,38 @@ def post_update_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for update_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_time_series` interceptor runs + before the `post_update_tensorboard_time_series_with_metadata` interceptor. """ return response + def post_update_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_update_tensorboard_time_series_with_metadata` interceptor runs after the + `post_update_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_time_series` will be passed to + `post_update_tensorboard_time_series_with_metadata`. + """ + return response, metadata + def pre_write_tensorboard_experiment_data( self, request: tensorboard_service.WriteTensorboardExperimentDataRequest, @@ -1084,12 +1785,38 @@ def post_write_tensorboard_experiment_data( ) -> tensorboard_service.WriteTensorboardExperimentDataResponse: """Post-rpc interceptor for write_tensorboard_experiment_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_experiment_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_experiment_data` interceptor runs + before the `post_write_tensorboard_experiment_data_with_metadata` interceptor. """ return response + def post_write_tensorboard_experiment_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardExperimentDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardExperimentDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_experiment_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_experiment_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_experiment_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_experiment_data_with_metadata` interceptor runs after the + `post_write_tensorboard_experiment_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_experiment_data` will be passed to + `post_write_tensorboard_experiment_data_with_metadata`. + """ + return response, metadata + def pre_write_tensorboard_run_data( self, request: tensorboard_service.WriteTensorboardRunDataRequest, @@ -1110,12 +1837,38 @@ def post_write_tensorboard_run_data( ) -> tensorboard_service.WriteTensorboardRunDataResponse: """Post-rpc interceptor for write_tensorboard_run_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_run_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_run_data` interceptor runs + before the `post_write_tensorboard_run_data_with_metadata` interceptor. """ return response + def post_write_tensorboard_run_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardRunDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardRunDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_run_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_run_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_run_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_run_data_with_metadata` interceptor runs after the + `post_write_tensorboard_run_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_run_data` will be passed to + `post_write_tensorboard_run_data_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3664,6 +4417,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_create_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3829,6 +4589,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3986,6 +4753,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_batch_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_batch_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4140,6 +4914,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4296,6 +5074,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4457,6 +5242,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4613,6 +5402,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4763,6 +5559,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4910,6 +5710,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5058,6 +5865,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5205,6 +6016,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_delete_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5366,6 +6184,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_export_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_export_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5517,6 +6342,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5669,6 +6498,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5820,6 +6653,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5970,6 +6807,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6120,6 +6961,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_experiments(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_experiments_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6273,6 +7118,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6424,6 +7273,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboards(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboards_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6574,6 +7427,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6728,6 +7585,10 @@ def __call__( ) resp = self._interceptor.post_read_tensorboard_blob_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_blob_data_with_metadata( + resp, response_metadata + ) return resp class _ReadTensorboardSize( @@ -6856,6 +7717,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_size(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_size_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7008,6 +7873,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7159,6 +8031,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_read_tensorboard_usage(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_read_tensorboard_usage_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7315,6 +8191,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7471,6 +8351,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7632,6 +8519,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7788,6 +8679,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_update_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -7950,6 +8848,13 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_tensorboard_experiment_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_write_tensorboard_experiment_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -8108,6 +9013,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_write_tensorboard_run_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_write_tensorboard_run_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest_asyncio.py index 5721e422d0..8f72ad6cfe 100644 --- a/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/tensorboard_service/transports/rest_asyncio.py @@ -374,12 +374,38 @@ async def post_batch_create_tensorboard_runs( ) -> tensorboard_service.BatchCreateTensorboardRunsResponse: """Post-rpc interceptor for batch_create_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_runs` interceptor runs + before the `post_batch_create_tensorboard_runs_with_metadata` interceptor. """ return response + async def post_batch_create_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_runs` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_runs_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_runs` will be passed to + `post_batch_create_tensorboard_runs_with_metadata`. + """ + return response, metadata + async def pre_batch_create_tensorboard_time_series( self, request: tensorboard_service.BatchCreateTensorboardTimeSeriesRequest, @@ -400,12 +426,38 @@ async def post_batch_create_tensorboard_time_series( ) -> tensorboard_service.BatchCreateTensorboardTimeSeriesResponse: """Post-rpc interceptor for batch_create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_create_tensorboard_time_series` interceptor runs + before the `post_batch_create_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_batch_create_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_batch_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_batch_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_batch_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_batch_create_tensorboard_time_series` will be passed to + `post_batch_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_batch_read_tensorboard_time_series_data( self, request: tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest, @@ -426,12 +478,38 @@ async def post_batch_read_tensorboard_time_series_data( ) -> tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for batch_read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_batch_read_tensorboard_time_series_data` interceptor runs + before the `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_batch_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for batch_read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_batch_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_batch_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_batch_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_batch_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_batch_read_tensorboard_time_series_data` will be passed to + `post_batch_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard( self, request: tensorboard_service.CreateTensorboardRequest, @@ -452,12 +530,35 @@ async def post_create_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard` interceptor runs + before the `post_create_tensorboard_with_metadata` interceptor. """ return response + async def post_create_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_with_metadata` + interceptor in new development instead of the `post_create_tensorboard` interceptor. + When both interceptors are used, this `post_create_tensorboard_with_metadata` interceptor runs after the + `post_create_tensorboard` interceptor. The (possibly modified) response returned by + `post_create_tensorboard` will be passed to + `post_create_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_experiment( self, request: tensorboard_service.CreateTensorboardExperimentRequest, @@ -478,12 +579,38 @@ async def post_create_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for create_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_experiment` interceptor runs + before the `post_create_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_create_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_create_tensorboard_experiment_with_metadata` interceptor runs after the + `post_create_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_experiment` will be passed to + `post_create_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_run( self, request: tensorboard_service.CreateTensorboardRunRequest, @@ -504,12 +631,37 @@ async def post_create_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for create_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_run` interceptor runs + before the `post_create_tensorboard_run_with_metadata` interceptor. """ return response + async def post_create_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for create_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_run` interceptor. + When both interceptors are used, this `post_create_tensorboard_run_with_metadata` interceptor runs after the + `post_create_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_run` will be passed to + `post_create_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_create_tensorboard_time_series( self, request: tensorboard_service.CreateTensorboardTimeSeriesRequest, @@ -530,12 +682,38 @@ async def post_create_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for create_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_create_tensorboard_time_series` interceptor runs + before the `post_create_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_create_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for create_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_create_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_create_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_create_tensorboard_time_series_with_metadata` interceptor runs after the + `post_create_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_create_tensorboard_time_series` will be passed to + `post_create_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard( self, request: tensorboard_service.DeleteTensorboardRequest, @@ -556,12 +734,35 @@ async def post_delete_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard` interceptor runs + before the `post_delete_tensorboard_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard` interceptor. + When both interceptors are used, this `post_delete_tensorboard_with_metadata` interceptor runs after the + `post_delete_tensorboard` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard` will be passed to + `post_delete_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_experiment( self, request: tensorboard_service.DeleteTensorboardExperimentRequest, @@ -582,12 +783,35 @@ async def post_delete_tensorboard_experiment( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_experiment` interceptor runs + before the `post_delete_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_experiment_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_delete_tensorboard_experiment_with_metadata` interceptor runs after the + `post_delete_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_experiment` will be passed to + `post_delete_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_run( self, request: tensorboard_service.DeleteTensorboardRunRequest, @@ -608,12 +832,35 @@ async def post_delete_tensorboard_run( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_run` interceptor runs + before the `post_delete_tensorboard_run_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_run_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_run` interceptor. + When both interceptors are used, this `post_delete_tensorboard_run_with_metadata` interceptor runs after the + `post_delete_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_run` will be passed to + `post_delete_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_delete_tensorboard_time_series( self, request: tensorboard_service.DeleteTensorboardTimeSeriesRequest, @@ -634,12 +881,35 @@ async def post_delete_tensorboard_time_series( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_delete_tensorboard_time_series` interceptor runs + before the `post_delete_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_delete_tensorboard_time_series_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_delete_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_delete_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_delete_tensorboard_time_series_with_metadata` interceptor runs after the + `post_delete_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_delete_tensorboard_time_series` will be passed to + `post_delete_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_export_tensorboard_time_series_data( self, request: tensorboard_service.ExportTensorboardTimeSeriesDataRequest, @@ -660,12 +930,38 @@ async def post_export_tensorboard_time_series_data( ) -> tensorboard_service.ExportTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for export_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_export_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_export_tensorboard_time_series_data` interceptor runs + before the `post_export_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_export_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ExportTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for export_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_export_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_export_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_export_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_export_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_export_tensorboard_time_series_data` will be passed to + `post_export_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard( self, request: tensorboard_service.GetTensorboardRequest, @@ -686,12 +982,35 @@ async def post_get_tensorboard( ) -> tensorboard.Tensorboard: """Post-rpc interceptor for get_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard` interceptor runs + before the `post_get_tensorboard_with_metadata` interceptor. """ return response + async def post_get_tensorboard_with_metadata( + self, + response: tensorboard.Tensorboard, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard.Tensorboard, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_with_metadata` + interceptor in new development instead of the `post_get_tensorboard` interceptor. + When both interceptors are used, this `post_get_tensorboard_with_metadata` interceptor runs after the + `post_get_tensorboard` interceptor. The (possibly modified) response returned by + `post_get_tensorboard` will be passed to + `post_get_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_experiment( self, request: tensorboard_service.GetTensorboardExperimentRequest, @@ -712,12 +1031,38 @@ async def post_get_tensorboard_experiment( ) -> tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for get_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_experiment` interceptor runs + before the `post_get_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_get_tensorboard_experiment_with_metadata( + self, + response: tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_get_tensorboard_experiment_with_metadata` interceptor runs after the + `post_get_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_experiment` will be passed to + `post_get_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_run( self, request: tensorboard_service.GetTensorboardRunRequest, @@ -738,12 +1083,35 @@ async def post_get_tensorboard_run( ) -> tensorboard_run.TensorboardRun: """Post-rpc interceptor for get_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_run` interceptor runs + before the `post_get_tensorboard_run_with_metadata` interceptor. """ return response + async def post_get_tensorboard_run_with_metadata( + self, + response: tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_run` interceptor. + When both interceptors are used, this `post_get_tensorboard_run_with_metadata` interceptor runs after the + `post_get_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_run` will be passed to + `post_get_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_get_tensorboard_time_series( self, request: tensorboard_service.GetTensorboardTimeSeriesRequest, @@ -764,12 +1132,38 @@ async def post_get_tensorboard_time_series( ) -> tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for get_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_get_tensorboard_time_series` interceptor runs + before the `post_get_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_get_tensorboard_time_series_with_metadata( + self, + response: tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for get_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_get_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_get_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_get_tensorboard_time_series_with_metadata` interceptor runs after the + `post_get_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_get_tensorboard_time_series` will be passed to + `post_get_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_experiments( self, request: tensorboard_service.ListTensorboardExperimentsRequest, @@ -790,12 +1184,38 @@ async def post_list_tensorboard_experiments( ) -> tensorboard_service.ListTensorboardExperimentsResponse: """Post-rpc interceptor for list_tensorboard_experiments - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_experiments_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_experiments` interceptor runs + before the `post_list_tensorboard_experiments_with_metadata` interceptor. """ return response + async def post_list_tensorboard_experiments_with_metadata( + self, + response: tensorboard_service.ListTensorboardExperimentsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardExperimentsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_experiments + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_experiments_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_experiments` interceptor. + When both interceptors are used, this `post_list_tensorboard_experiments_with_metadata` interceptor runs after the + `post_list_tensorboard_experiments` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_experiments` will be passed to + `post_list_tensorboard_experiments_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_runs( self, request: tensorboard_service.ListTensorboardRunsRequest, @@ -816,12 +1236,38 @@ async def post_list_tensorboard_runs( ) -> tensorboard_service.ListTensorboardRunsResponse: """Post-rpc interceptor for list_tensorboard_runs - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_runs_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_runs` interceptor runs + before the `post_list_tensorboard_runs_with_metadata` interceptor. """ return response + async def post_list_tensorboard_runs_with_metadata( + self, + response: tensorboard_service.ListTensorboardRunsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardRunsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_runs + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_runs_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_runs` interceptor. + When both interceptors are used, this `post_list_tensorboard_runs_with_metadata` interceptor runs after the + `post_list_tensorboard_runs` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_runs` will be passed to + `post_list_tensorboard_runs_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboards( self, request: tensorboard_service.ListTensorboardsRequest, @@ -842,12 +1288,38 @@ async def post_list_tensorboards( ) -> tensorboard_service.ListTensorboardsResponse: """Post-rpc interceptor for list_tensorboards - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboards_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboards` interceptor runs + before the `post_list_tensorboards_with_metadata` interceptor. """ return response + async def post_list_tensorboards_with_metadata( + self, + response: tensorboard_service.ListTensorboardsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboards + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboards_with_metadata` + interceptor in new development instead of the `post_list_tensorboards` interceptor. + When both interceptors are used, this `post_list_tensorboards_with_metadata` interceptor runs after the + `post_list_tensorboards` interceptor. The (possibly modified) response returned by + `post_list_tensorboards` will be passed to + `post_list_tensorboards_with_metadata`. + """ + return response, metadata + async def pre_list_tensorboard_time_series( self, request: tensorboard_service.ListTensorboardTimeSeriesRequest, @@ -868,12 +1340,38 @@ async def post_list_tensorboard_time_series( ) -> tensorboard_service.ListTensorboardTimeSeriesResponse: """Post-rpc interceptor for list_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_list_tensorboard_time_series` interceptor runs + before the `post_list_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_list_tensorboard_time_series_with_metadata( + self, + response: tensorboard_service.ListTensorboardTimeSeriesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ListTensorboardTimeSeriesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_list_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_list_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_list_tensorboard_time_series_with_metadata` interceptor runs after the + `post_list_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_list_tensorboard_time_series` will be passed to + `post_list_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_blob_data( self, request: tensorboard_service.ReadTensorboardBlobDataRequest, @@ -894,12 +1392,38 @@ async def post_read_tensorboard_blob_data( ) -> rest_streaming_async.AsyncResponseIterator: """Post-rpc interceptor for read_tensorboard_blob_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_blob_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_blob_data` interceptor runs + before the `post_read_tensorboard_blob_data_with_metadata` interceptor. """ return response + async def post_read_tensorboard_blob_data_with_metadata( + self, + response: rest_streaming_async.AsyncResponseIterator, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + rest_streaming_async.AsyncResponseIterator, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_blob_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_blob_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_blob_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_blob_data_with_metadata` interceptor runs after the + `post_read_tensorboard_blob_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_blob_data` will be passed to + `post_read_tensorboard_blob_data_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_size( self, request: tensorboard_service.ReadTensorboardSizeRequest, @@ -920,12 +1444,38 @@ async def post_read_tensorboard_size( ) -> tensorboard_service.ReadTensorboardSizeResponse: """Post-rpc interceptor for read_tensorboard_size - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_size_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_size` interceptor runs + before the `post_read_tensorboard_size_with_metadata` interceptor. """ return response + async def post_read_tensorboard_size_with_metadata( + self, + response: tensorboard_service.ReadTensorboardSizeResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardSizeResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_size + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_size_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_size` interceptor. + When both interceptors are used, this `post_read_tensorboard_size_with_metadata` interceptor runs after the + `post_read_tensorboard_size` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_size` will be passed to + `post_read_tensorboard_size_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_time_series_data( self, request: tensorboard_service.ReadTensorboardTimeSeriesDataRequest, @@ -946,12 +1496,38 @@ async def post_read_tensorboard_time_series_data( ) -> tensorboard_service.ReadTensorboardTimeSeriesDataResponse: """Post-rpc interceptor for read_tensorboard_time_series_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_time_series_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_time_series_data` interceptor runs + before the `post_read_tensorboard_time_series_data_with_metadata` interceptor. """ return response + async def post_read_tensorboard_time_series_data_with_metadata( + self, + response: tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardTimeSeriesDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_time_series_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_time_series_data_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_time_series_data` interceptor. + When both interceptors are used, this `post_read_tensorboard_time_series_data_with_metadata` interceptor runs after the + `post_read_tensorboard_time_series_data` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_time_series_data` will be passed to + `post_read_tensorboard_time_series_data_with_metadata`. + """ + return response, metadata + async def pre_read_tensorboard_usage( self, request: tensorboard_service.ReadTensorboardUsageRequest, @@ -972,12 +1548,38 @@ async def post_read_tensorboard_usage( ) -> tensorboard_service.ReadTensorboardUsageResponse: """Post-rpc interceptor for read_tensorboard_usage - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_read_tensorboard_usage_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_read_tensorboard_usage` interceptor runs + before the `post_read_tensorboard_usage_with_metadata` interceptor. """ return response + async def post_read_tensorboard_usage_with_metadata( + self, + response: tensorboard_service.ReadTensorboardUsageResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.ReadTensorboardUsageResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for read_tensorboard_usage + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_read_tensorboard_usage_with_metadata` + interceptor in new development instead of the `post_read_tensorboard_usage` interceptor. + When both interceptors are used, this `post_read_tensorboard_usage_with_metadata` interceptor runs after the + `post_read_tensorboard_usage` interceptor. The (possibly modified) response returned by + `post_read_tensorboard_usage` will be passed to + `post_read_tensorboard_usage_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard( self, request: tensorboard_service.UpdateTensorboardRequest, @@ -998,12 +1600,35 @@ async def post_update_tensorboard( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_tensorboard - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard` interceptor runs + before the `post_update_tensorboard_with_metadata` interceptor. """ return response + async def post_update_tensorboard_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_tensorboard + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_with_metadata` + interceptor in new development instead of the `post_update_tensorboard` interceptor. + When both interceptors are used, this `post_update_tensorboard_with_metadata` interceptor runs after the + `post_update_tensorboard` interceptor. The (possibly modified) response returned by + `post_update_tensorboard` will be passed to + `post_update_tensorboard_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_experiment( self, request: tensorboard_service.UpdateTensorboardExperimentRequest, @@ -1024,12 +1649,38 @@ async def post_update_tensorboard_experiment( ) -> gca_tensorboard_experiment.TensorboardExperiment: """Post-rpc interceptor for update_tensorboard_experiment - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_experiment_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_experiment` interceptor runs + before the `post_update_tensorboard_experiment_with_metadata` interceptor. """ return response + async def post_update_tensorboard_experiment_with_metadata( + self, + response: gca_tensorboard_experiment.TensorboardExperiment, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_experiment.TensorboardExperiment, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_experiment + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_experiment_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_experiment` interceptor. + When both interceptors are used, this `post_update_tensorboard_experiment_with_metadata` interceptor runs after the + `post_update_tensorboard_experiment` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_experiment` will be passed to + `post_update_tensorboard_experiment_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_run( self, request: tensorboard_service.UpdateTensorboardRunRequest, @@ -1050,12 +1701,37 @@ async def post_update_tensorboard_run( ) -> gca_tensorboard_run.TensorboardRun: """Post-rpc interceptor for update_tensorboard_run - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_run_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_run` interceptor runs + before the `post_update_tensorboard_run_with_metadata` interceptor. """ return response + async def post_update_tensorboard_run_with_metadata( + self, + response: gca_tensorboard_run.TensorboardRun, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_run.TensorboardRun, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for update_tensorboard_run + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_run_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_run` interceptor. + When both interceptors are used, this `post_update_tensorboard_run_with_metadata` interceptor runs after the + `post_update_tensorboard_run` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_run` will be passed to + `post_update_tensorboard_run_with_metadata`. + """ + return response, metadata + async def pre_update_tensorboard_time_series( self, request: tensorboard_service.UpdateTensorboardTimeSeriesRequest, @@ -1076,12 +1752,38 @@ async def post_update_tensorboard_time_series( ) -> gca_tensorboard_time_series.TensorboardTimeSeries: """Post-rpc interceptor for update_tensorboard_time_series - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_tensorboard_time_series_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_update_tensorboard_time_series` interceptor runs + before the `post_update_tensorboard_time_series_with_metadata` interceptor. """ return response + async def post_update_tensorboard_time_series_with_metadata( + self, + response: gca_tensorboard_time_series.TensorboardTimeSeries, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + gca_tensorboard_time_series.TensorboardTimeSeries, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for update_tensorboard_time_series + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_update_tensorboard_time_series_with_metadata` + interceptor in new development instead of the `post_update_tensorboard_time_series` interceptor. + When both interceptors are used, this `post_update_tensorboard_time_series_with_metadata` interceptor runs after the + `post_update_tensorboard_time_series` interceptor. The (possibly modified) response returned by + `post_update_tensorboard_time_series` will be passed to + `post_update_tensorboard_time_series_with_metadata`. + """ + return response, metadata + async def pre_write_tensorboard_experiment_data( self, request: tensorboard_service.WriteTensorboardExperimentDataRequest, @@ -1102,12 +1804,38 @@ async def post_write_tensorboard_experiment_data( ) -> tensorboard_service.WriteTensorboardExperimentDataResponse: """Post-rpc interceptor for write_tensorboard_experiment_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_experiment_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_experiment_data` interceptor runs + before the `post_write_tensorboard_experiment_data_with_metadata` interceptor. """ return response + async def post_write_tensorboard_experiment_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardExperimentDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardExperimentDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_experiment_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_experiment_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_experiment_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_experiment_data_with_metadata` interceptor runs after the + `post_write_tensorboard_experiment_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_experiment_data` will be passed to + `post_write_tensorboard_experiment_data_with_metadata`. + """ + return response, metadata + async def pre_write_tensorboard_run_data( self, request: tensorboard_service.WriteTensorboardRunDataRequest, @@ -1128,12 +1856,38 @@ async def post_write_tensorboard_run_data( ) -> tensorboard_service.WriteTensorboardRunDataResponse: """Post-rpc interceptor for write_tensorboard_run_data - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_write_tensorboard_run_data_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the TensorboardService server but before - it is returned to user code. + it is returned to user code. This `post_write_tensorboard_run_data` interceptor runs + before the `post_write_tensorboard_run_data_with_metadata` interceptor. """ return response + async def post_write_tensorboard_run_data_with_metadata( + self, + response: tensorboard_service.WriteTensorboardRunDataResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + tensorboard_service.WriteTensorboardRunDataResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for write_tensorboard_run_data + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the TensorboardService server but before it is returned to user code. + + We recommend only using this `post_write_tensorboard_run_data_with_metadata` + interceptor in new development instead of the `post_write_tensorboard_run_data` interceptor. + When both interceptors are used, this `post_write_tensorboard_run_data_with_metadata` interceptor runs after the + `post_write_tensorboard_run_data` interceptor. The (possibly modified) response returned by + `post_write_tensorboard_run_data` will be passed to + `post_write_tensorboard_run_data_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1797,6 +2551,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_batch_create_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_create_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1970,6 +2731,13 @@ async def __call__( resp = await self._interceptor.post_batch_create_tensorboard_time_series( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2135,6 +2903,13 @@ async def __call__( resp = await self._interceptor.post_batch_read_tensorboard_time_series_data( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_batch_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2297,6 +3072,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2464,6 +3243,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2629,6 +3415,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2796,6 +3586,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_create_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2954,6 +3751,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3114,6 +3915,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3268,6 +4076,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3428,6 +4240,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_delete_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3597,6 +4416,13 @@ async def __call__( resp = await self._interceptor.post_export_tensorboard_time_series_data( resp ) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_export_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3756,6 +4582,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3912,6 +4742,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4069,6 +4906,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4223,6 +5064,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_get_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4384,6 +5232,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_experiments(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_tensorboard_experiments_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4541,6 +5396,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_runs(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tensorboard_runs_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4698,6 +5557,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboards(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_tensorboards_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4859,6 +5722,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_list_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5016,6 +5886,13 @@ async def __call__( response, tensorboard_service.ReadTensorboardBlobDataResponse ) resp = await self._interceptor.post_read_tensorboard_blob_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_read_tensorboard_blob_data_with_metadata( + resp, response_metadata + ) return resp class _ReadTensorboardSize( @@ -5147,6 +6024,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_size(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_tensorboard_size_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5310,6 +6191,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_time_series_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_read_tensorboard_time_series_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5465,6 +6353,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_read_tensorboard_usage(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_read_tensorboard_usage_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5629,6 +6521,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_tensorboard_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5796,6 +6692,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_experiment(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_tensorboard_experiment_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5961,6 +6864,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_run(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_tensorboard_run_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6128,6 +7035,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_tensorboard_time_series(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_update_tensorboard_time_series_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6299,6 +7213,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_tensorboard_experiment_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_write_tensorboard_experiment_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -6461,6 +7382,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_write_tensorboard_run_data(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_write_tensorboard_run_data_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/client.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/client.py index 5ecbc06b7b..3a5eeab0b1 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -599,6 +601,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2134,16 +2163,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2189,16 +2222,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2361,16 +2398,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2482,16 +2523,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2604,16 +2649,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2664,16 +2713,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2719,16 +2772,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -2774,16 +2831,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest.py index 7425f02077..3b37b8cfe5 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest.py @@ -186,12 +186,35 @@ def post_create_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_create_rag_corpus` interceptor runs + before the `post_create_rag_corpus_with_metadata` interceptor. """ return response + def post_create_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_create_rag_corpus_with_metadata` + interceptor in new development instead of the `post_create_rag_corpus` interceptor. + When both interceptors are used, this `post_create_rag_corpus_with_metadata` interceptor runs after the + `post_create_rag_corpus` interceptor. The (possibly modified) response returned by + `post_create_rag_corpus` will be passed to + `post_create_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_delete_rag_corpus( self, request: vertex_rag_data_service.DeleteRagCorpusRequest, @@ -212,12 +235,35 @@ def post_delete_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_corpus` interceptor runs + before the `post_delete_rag_corpus_with_metadata` interceptor. """ return response + def post_delete_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_corpus_with_metadata` + interceptor in new development instead of the `post_delete_rag_corpus` interceptor. + When both interceptors are used, this `post_delete_rag_corpus_with_metadata` interceptor runs after the + `post_delete_rag_corpus` interceptor. The (possibly modified) response returned by + `post_delete_rag_corpus` will be passed to + `post_delete_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_delete_rag_file( self, request: vertex_rag_data_service.DeleteRagFileRequest, @@ -238,12 +284,35 @@ def post_delete_rag_file( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_file` interceptor runs + before the `post_delete_rag_file_with_metadata` interceptor. """ return response + def post_delete_rag_file_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_file_with_metadata` + interceptor in new development instead of the `post_delete_rag_file` interceptor. + When both interceptors are used, this `post_delete_rag_file_with_metadata` interceptor runs after the + `post_delete_rag_file` interceptor. The (possibly modified) response returned by + `post_delete_rag_file` will be passed to + `post_delete_rag_file_with_metadata`. + """ + return response, metadata + def pre_get_rag_corpus( self, request: vertex_rag_data_service.GetRagCorpusRequest, @@ -264,12 +333,35 @@ def post_get_rag_corpus( ) -> vertex_rag_data.RagCorpus: """Post-rpc interceptor for get_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_corpus` interceptor runs + before the `post_get_rag_corpus_with_metadata` interceptor. """ return response + def post_get_rag_corpus_with_metadata( + self, + response: vertex_rag_data.RagCorpus, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagCorpus, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_corpus_with_metadata` + interceptor in new development instead of the `post_get_rag_corpus` interceptor. + When both interceptors are used, this `post_get_rag_corpus_with_metadata` interceptor runs after the + `post_get_rag_corpus` interceptor. The (possibly modified) response returned by + `post_get_rag_corpus` will be passed to + `post_get_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_get_rag_file( self, request: vertex_rag_data_service.GetRagFileRequest, @@ -290,12 +382,35 @@ def post_get_rag_file( ) -> vertex_rag_data.RagFile: """Post-rpc interceptor for get_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_file` interceptor runs + before the `post_get_rag_file_with_metadata` interceptor. """ return response + def post_get_rag_file_with_metadata( + self, + response: vertex_rag_data.RagFile, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagFile, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_file_with_metadata` + interceptor in new development instead of the `post_get_rag_file` interceptor. + When both interceptors are used, this `post_get_rag_file_with_metadata` interceptor runs after the + `post_get_rag_file` interceptor. The (possibly modified) response returned by + `post_get_rag_file` will be passed to + `post_get_rag_file_with_metadata`. + """ + return response, metadata + def pre_import_rag_files( self, request: vertex_rag_data_service.ImportRagFilesRequest, @@ -316,12 +431,35 @@ def post_import_rag_files( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_import_rag_files` interceptor runs + before the `post_import_rag_files_with_metadata` interceptor. """ return response + def post_import_rag_files_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_import_rag_files_with_metadata` + interceptor in new development instead of the `post_import_rag_files` interceptor. + When both interceptors are used, this `post_import_rag_files_with_metadata` interceptor runs after the + `post_import_rag_files` interceptor. The (possibly modified) response returned by + `post_import_rag_files` will be passed to + `post_import_rag_files_with_metadata`. + """ + return response, metadata + def pre_list_rag_corpora( self, request: vertex_rag_data_service.ListRagCorporaRequest, @@ -342,12 +480,38 @@ def post_list_rag_corpora( ) -> vertex_rag_data_service.ListRagCorporaResponse: """Post-rpc interceptor for list_rag_corpora - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_corpora_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_corpora` interceptor runs + before the `post_list_rag_corpora_with_metadata` interceptor. """ return response + def post_list_rag_corpora_with_metadata( + self, + response: vertex_rag_data_service.ListRagCorporaResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagCorporaResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_corpora + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_corpora_with_metadata` + interceptor in new development instead of the `post_list_rag_corpora` interceptor. + When both interceptors are used, this `post_list_rag_corpora_with_metadata` interceptor runs after the + `post_list_rag_corpora` interceptor. The (possibly modified) response returned by + `post_list_rag_corpora` will be passed to + `post_list_rag_corpora_with_metadata`. + """ + return response, metadata + def pre_list_rag_files( self, request: vertex_rag_data_service.ListRagFilesRequest, @@ -368,12 +532,38 @@ def post_list_rag_files( ) -> vertex_rag_data_service.ListRagFilesResponse: """Post-rpc interceptor for list_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_files` interceptor runs + before the `post_list_rag_files_with_metadata` interceptor. """ return response + def post_list_rag_files_with_metadata( + self, + response: vertex_rag_data_service.ListRagFilesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagFilesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_files_with_metadata` + interceptor in new development instead of the `post_list_rag_files` interceptor. + When both interceptors are used, this `post_list_rag_files_with_metadata` interceptor runs after the + `post_list_rag_files` interceptor. The (possibly modified) response returned by + `post_list_rag_files` will be passed to + `post_list_rag_files_with_metadata`. + """ + return response, metadata + def pre_update_rag_corpus( self, request: vertex_rag_data_service.UpdateRagCorpusRequest, @@ -394,12 +584,35 @@ def post_update_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_update_rag_corpus` interceptor runs + before the `post_update_rag_corpus_with_metadata` interceptor. """ return response + def post_update_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_update_rag_corpus_with_metadata` + interceptor in new development instead of the `post_update_rag_corpus` interceptor. + When both interceptors are used, this `post_update_rag_corpus_with_metadata` interceptor runs after the + `post_update_rag_corpus` interceptor. The (possibly modified) response returned by + `post_update_rag_corpus` will be passed to + `post_update_rag_corpus_with_metadata`. + """ + return response, metadata + def pre_upload_rag_file( self, request: vertex_rag_data_service.UploadRagFileRequest, @@ -420,12 +633,38 @@ def post_upload_rag_file( ) -> vertex_rag_data_service.UploadRagFileResponse: """Post-rpc interceptor for upload_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_upload_rag_file` interceptor runs + before the `post_upload_rag_file_with_metadata` interceptor. """ return response + def post_upload_rag_file_with_metadata( + self, + response: vertex_rag_data_service.UploadRagFileResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.UploadRagFileResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for upload_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_upload_rag_file_with_metadata` + interceptor in new development instead of the `post_upload_rag_file` interceptor. + When both interceptors are used, this `post_upload_rag_file_with_metadata` interceptor runs after the + `post_upload_rag_file` interceptor. The (possibly modified) response returned by + `post_upload_rag_file` will be passed to + `post_upload_rag_file_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -2972,6 +3211,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3118,6 +3361,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3262,6 +3509,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_delete_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_delete_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3408,6 +3659,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3553,6 +3808,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3705,6 +3964,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_import_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_import_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3852,6 +4115,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_rag_corpora(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_rag_corpora_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3999,6 +4266,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4153,6 +4424,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_update_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_update_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4304,6 +4579,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_upload_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_upload_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest_asyncio.py index c4702b6d4e..984b1fa3b5 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_data_service/transports/rest_asyncio.py @@ -204,12 +204,35 @@ async def post_create_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for create_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_create_rag_corpus` interceptor runs + before the `post_create_rag_corpus_with_metadata` interceptor. """ return response + async def post_create_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_create_rag_corpus_with_metadata` + interceptor in new development instead of the `post_create_rag_corpus` interceptor. + When both interceptors are used, this `post_create_rag_corpus_with_metadata` interceptor runs after the + `post_create_rag_corpus` interceptor. The (possibly modified) response returned by + `post_create_rag_corpus` will be passed to + `post_create_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_delete_rag_corpus( self, request: vertex_rag_data_service.DeleteRagCorpusRequest, @@ -230,12 +253,35 @@ async def post_delete_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_corpus` interceptor runs + before the `post_delete_rag_corpus_with_metadata` interceptor. """ return response + async def post_delete_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_corpus_with_metadata` + interceptor in new development instead of the `post_delete_rag_corpus` interceptor. + When both interceptors are used, this `post_delete_rag_corpus_with_metadata` interceptor runs after the + `post_delete_rag_corpus` interceptor. The (possibly modified) response returned by + `post_delete_rag_corpus` will be passed to + `post_delete_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_delete_rag_file( self, request: vertex_rag_data_service.DeleteRagFileRequest, @@ -256,12 +302,35 @@ async def post_delete_rag_file( ) -> operations_pb2.Operation: """Post-rpc interceptor for delete_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_delete_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_delete_rag_file` interceptor runs + before the `post_delete_rag_file_with_metadata` interceptor. """ return response + async def post_delete_rag_file_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for delete_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_delete_rag_file_with_metadata` + interceptor in new development instead of the `post_delete_rag_file` interceptor. + When both interceptors are used, this `post_delete_rag_file_with_metadata` interceptor runs after the + `post_delete_rag_file` interceptor. The (possibly modified) response returned by + `post_delete_rag_file` will be passed to + `post_delete_rag_file_with_metadata`. + """ + return response, metadata + async def pre_get_rag_corpus( self, request: vertex_rag_data_service.GetRagCorpusRequest, @@ -282,12 +351,35 @@ async def post_get_rag_corpus( ) -> vertex_rag_data.RagCorpus: """Post-rpc interceptor for get_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_corpus` interceptor runs + before the `post_get_rag_corpus_with_metadata` interceptor. """ return response + async def post_get_rag_corpus_with_metadata( + self, + response: vertex_rag_data.RagCorpus, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagCorpus, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_corpus_with_metadata` + interceptor in new development instead of the `post_get_rag_corpus` interceptor. + When both interceptors are used, this `post_get_rag_corpus_with_metadata` interceptor runs after the + `post_get_rag_corpus` interceptor. The (possibly modified) response returned by + `post_get_rag_corpus` will be passed to + `post_get_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_get_rag_file( self, request: vertex_rag_data_service.GetRagFileRequest, @@ -308,12 +400,35 @@ async def post_get_rag_file( ) -> vertex_rag_data.RagFile: """Post-rpc interceptor for get_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_get_rag_file` interceptor runs + before the `post_get_rag_file_with_metadata` interceptor. """ return response + async def post_get_rag_file_with_metadata( + self, + response: vertex_rag_data.RagFile, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[vertex_rag_data.RagFile, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_get_rag_file_with_metadata` + interceptor in new development instead of the `post_get_rag_file` interceptor. + When both interceptors are used, this `post_get_rag_file_with_metadata` interceptor runs after the + `post_get_rag_file` interceptor. The (possibly modified) response returned by + `post_get_rag_file` will be passed to + `post_get_rag_file_with_metadata`. + """ + return response, metadata + async def pre_import_rag_files( self, request: vertex_rag_data_service.ImportRagFilesRequest, @@ -334,12 +449,35 @@ async def post_import_rag_files( ) -> operations_pb2.Operation: """Post-rpc interceptor for import_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_import_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_import_rag_files` interceptor runs + before the `post_import_rag_files_with_metadata` interceptor. """ return response + async def post_import_rag_files_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for import_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_import_rag_files_with_metadata` + interceptor in new development instead of the `post_import_rag_files` interceptor. + When both interceptors are used, this `post_import_rag_files_with_metadata` interceptor runs after the + `post_import_rag_files` interceptor. The (possibly modified) response returned by + `post_import_rag_files` will be passed to + `post_import_rag_files_with_metadata`. + """ + return response, metadata + async def pre_list_rag_corpora( self, request: vertex_rag_data_service.ListRagCorporaRequest, @@ -360,12 +498,38 @@ async def post_list_rag_corpora( ) -> vertex_rag_data_service.ListRagCorporaResponse: """Post-rpc interceptor for list_rag_corpora - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_corpora_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_corpora` interceptor runs + before the `post_list_rag_corpora_with_metadata` interceptor. """ return response + async def post_list_rag_corpora_with_metadata( + self, + response: vertex_rag_data_service.ListRagCorporaResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagCorporaResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_corpora + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_corpora_with_metadata` + interceptor in new development instead of the `post_list_rag_corpora` interceptor. + When both interceptors are used, this `post_list_rag_corpora_with_metadata` interceptor runs after the + `post_list_rag_corpora` interceptor. The (possibly modified) response returned by + `post_list_rag_corpora` will be passed to + `post_list_rag_corpora_with_metadata`. + """ + return response, metadata + async def pre_list_rag_files( self, request: vertex_rag_data_service.ListRagFilesRequest, @@ -386,12 +550,38 @@ async def post_list_rag_files( ) -> vertex_rag_data_service.ListRagFilesResponse: """Post-rpc interceptor for list_rag_files - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_rag_files_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_list_rag_files` interceptor runs + before the `post_list_rag_files_with_metadata` interceptor. """ return response + async def post_list_rag_files_with_metadata( + self, + response: vertex_rag_data_service.ListRagFilesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.ListRagFilesResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_rag_files + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_list_rag_files_with_metadata` + interceptor in new development instead of the `post_list_rag_files` interceptor. + When both interceptors are used, this `post_list_rag_files_with_metadata` interceptor runs after the + `post_list_rag_files` interceptor. The (possibly modified) response returned by + `post_list_rag_files` will be passed to + `post_list_rag_files_with_metadata`. + """ + return response, metadata + async def pre_update_rag_corpus( self, request: vertex_rag_data_service.UpdateRagCorpusRequest, @@ -412,12 +602,35 @@ async def post_update_rag_corpus( ) -> operations_pb2.Operation: """Post-rpc interceptor for update_rag_corpus - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_update_rag_corpus_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_update_rag_corpus` interceptor runs + before the `post_update_rag_corpus_with_metadata` interceptor. """ return response + async def post_update_rag_corpus_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for update_rag_corpus + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_update_rag_corpus_with_metadata` + interceptor in new development instead of the `post_update_rag_corpus` interceptor. + When both interceptors are used, this `post_update_rag_corpus_with_metadata` interceptor runs after the + `post_update_rag_corpus` interceptor. The (possibly modified) response returned by + `post_update_rag_corpus` will be passed to + `post_update_rag_corpus_with_metadata`. + """ + return response, metadata + async def pre_upload_rag_file( self, request: vertex_rag_data_service.UploadRagFileRequest, @@ -438,12 +651,38 @@ async def post_upload_rag_file( ) -> vertex_rag_data_service.UploadRagFileResponse: """Post-rpc interceptor for upload_rag_file - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_upload_rag_file_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagDataService server but before - it is returned to user code. + it is returned to user code. This `post_upload_rag_file` interceptor runs + before the `post_upload_rag_file_with_metadata` interceptor. """ return response + async def post_upload_rag_file_with_metadata( + self, + response: vertex_rag_data_service.UploadRagFileResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_data_service.UploadRagFileResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for upload_rag_file + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagDataService server but before it is returned to user code. + + We recommend only using this `post_upload_rag_file_with_metadata` + interceptor in new development instead of the `post_upload_rag_file` interceptor. + When both interceptors are used, this `post_upload_rag_file_with_metadata` interceptor runs after the + `post_upload_rag_file` interceptor. The (possibly modified) response returned by + `post_upload_rag_file` will be passed to + `post_upload_rag_file_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1002,6 +1241,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1156,6 +1399,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1310,6 +1557,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_delete_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_delete_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1464,6 +1715,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1619,6 +1874,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1779,6 +2038,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_import_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_import_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1932,6 +2195,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_rag_corpora(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_rag_corpora_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2087,6 +2354,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_rag_files(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_rag_files_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2249,6 +2520,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_update_rag_corpus(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_update_rag_corpus_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2408,6 +2683,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_upload_rag_file(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_upload_rag_file_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/client.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/client.py index 25fc1040ef..41e6ba8659 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -503,6 +505,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -1165,16 +1194,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -1220,16 +1253,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -1392,16 +1429,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -1513,16 +1554,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -1635,16 +1680,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -1695,16 +1744,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -1750,16 +1803,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -1805,16 +1862,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest.py index c552f8ebfe..536ef6a752 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest.py @@ -127,12 +127,38 @@ def post_augment_prompt( ) -> vertex_rag_service.AugmentPromptResponse: """Post-rpc interceptor for augment_prompt - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_augment_prompt_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_augment_prompt` interceptor runs + before the `post_augment_prompt_with_metadata` interceptor. """ return response + def post_augment_prompt_with_metadata( + self, + response: vertex_rag_service.AugmentPromptResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.AugmentPromptResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for augment_prompt + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_augment_prompt_with_metadata` + interceptor in new development instead of the `post_augment_prompt` interceptor. + When both interceptors are used, this `post_augment_prompt_with_metadata` interceptor runs after the + `post_augment_prompt` interceptor. The (possibly modified) response returned by + `post_augment_prompt` will be passed to + `post_augment_prompt_with_metadata`. + """ + return response, metadata + def pre_corroborate_content( self, request: vertex_rag_service.CorroborateContentRequest, @@ -153,12 +179,38 @@ def post_corroborate_content( ) -> vertex_rag_service.CorroborateContentResponse: """Post-rpc interceptor for corroborate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_corroborate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_corroborate_content` interceptor runs + before the `post_corroborate_content_with_metadata` interceptor. """ return response + def post_corroborate_content_with_metadata( + self, + response: vertex_rag_service.CorroborateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.CorroborateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for corroborate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_corroborate_content_with_metadata` + interceptor in new development instead of the `post_corroborate_content` interceptor. + When both interceptors are used, this `post_corroborate_content_with_metadata` interceptor runs after the + `post_corroborate_content` interceptor. The (possibly modified) response returned by + `post_corroborate_content` will be passed to + `post_corroborate_content_with_metadata`. + """ + return response, metadata + def pre_retrieve_contexts( self, request: vertex_rag_service.RetrieveContextsRequest, @@ -179,12 +231,38 @@ def post_retrieve_contexts( ) -> vertex_rag_service.RetrieveContextsResponse: """Post-rpc interceptor for retrieve_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_retrieve_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_retrieve_contexts` interceptor runs + before the `post_retrieve_contexts_with_metadata` interceptor. """ return response + def post_retrieve_contexts_with_metadata( + self, + response: vertex_rag_service.RetrieveContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.RetrieveContextsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for retrieve_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_retrieve_contexts_with_metadata` + interceptor in new development instead of the `post_retrieve_contexts` interceptor. + When both interceptors are used, this `post_retrieve_contexts_with_metadata` interceptor runs after the + `post_retrieve_contexts` interceptor. The (possibly modified) response returned by + `post_retrieve_contexts` will be passed to + `post_retrieve_contexts_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -639,6 +717,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_augment_prompt(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_augment_prompt_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -794,6 +876,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_corroborate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_corroborate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -949,6 +1035,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_retrieve_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_retrieve_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest_asyncio.py index 8074b6578c..ac07dfdf0a 100644 --- a/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/vertex_rag_service/transports/rest_asyncio.py @@ -144,12 +144,38 @@ async def post_augment_prompt( ) -> vertex_rag_service.AugmentPromptResponse: """Post-rpc interceptor for augment_prompt - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_augment_prompt_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_augment_prompt` interceptor runs + before the `post_augment_prompt_with_metadata` interceptor. """ return response + async def post_augment_prompt_with_metadata( + self, + response: vertex_rag_service.AugmentPromptResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.AugmentPromptResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for augment_prompt + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_augment_prompt_with_metadata` + interceptor in new development instead of the `post_augment_prompt` interceptor. + When both interceptors are used, this `post_augment_prompt_with_metadata` interceptor runs after the + `post_augment_prompt` interceptor. The (possibly modified) response returned by + `post_augment_prompt` will be passed to + `post_augment_prompt_with_metadata`. + """ + return response, metadata + async def pre_corroborate_content( self, request: vertex_rag_service.CorroborateContentRequest, @@ -170,12 +196,38 @@ async def post_corroborate_content( ) -> vertex_rag_service.CorroborateContentResponse: """Post-rpc interceptor for corroborate_content - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_corroborate_content_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_corroborate_content` interceptor runs + before the `post_corroborate_content_with_metadata` interceptor. """ return response + async def post_corroborate_content_with_metadata( + self, + response: vertex_rag_service.CorroborateContentResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.CorroborateContentResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for corroborate_content + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_corroborate_content_with_metadata` + interceptor in new development instead of the `post_corroborate_content` interceptor. + When both interceptors are used, this `post_corroborate_content_with_metadata` interceptor runs after the + `post_corroborate_content` interceptor. The (possibly modified) response returned by + `post_corroborate_content` will be passed to + `post_corroborate_content_with_metadata`. + """ + return response, metadata + async def pre_retrieve_contexts( self, request: vertex_rag_service.RetrieveContextsRequest, @@ -196,12 +248,38 @@ async def post_retrieve_contexts( ) -> vertex_rag_service.RetrieveContextsResponse: """Post-rpc interceptor for retrieve_contexts - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_retrieve_contexts_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VertexRagService server but before - it is returned to user code. + it is returned to user code. This `post_retrieve_contexts` interceptor runs + before the `post_retrieve_contexts_with_metadata` interceptor. """ return response + async def post_retrieve_contexts_with_metadata( + self, + response: vertex_rag_service.RetrieveContextsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vertex_rag_service.RetrieveContextsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for retrieve_contexts + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VertexRagService server but before it is returned to user code. + + We recommend only using this `post_retrieve_contexts_with_metadata` + interceptor in new development instead of the `post_retrieve_contexts` interceptor. + When both interceptors are used, this `post_retrieve_contexts_with_metadata` interceptor runs after the + `post_retrieve_contexts` interceptor. The (possibly modified) response returned by + `post_retrieve_contexts` will be passed to + `post_retrieve_contexts_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -720,6 +798,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_augment_prompt(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_augment_prompt_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -881,6 +963,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_corroborate_content(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_corroborate_content_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1042,6 +1128,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_retrieve_contexts(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_retrieve_contexts_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vizier_service/client.py b/google/cloud/aiplatform_v1beta1/services/vizier_service/client.py index 5fc56b3bb8..197a01d34a 100644 --- a/google/cloud/aiplatform_v1beta1/services/vizier_service/client.py +++ b/google/cloud/aiplatform_v1beta1/services/vizier_service/client.py @@ -14,6 +14,8 @@ # limitations under the License. # from collections import OrderedDict +from http import HTTPStatus +import json import logging as std_logging import os import re @@ -555,6 +557,33 @@ def _validate_universe_domain(self): # NOTE (b/349488459): universe validation is disabled until further notice. return True + def _add_cred_info_for_auth_errors( + self, error: core_exceptions.GoogleAPICallError + ) -> None: + """Adds credential info string to error details for 401/403/404 errors. + + Args: + error (google.api_core.exceptions.GoogleAPICallError): The error to add the cred info. + """ + if error.code not in [ + HTTPStatus.UNAUTHORIZED, + HTTPStatus.FORBIDDEN, + HTTPStatus.NOT_FOUND, + ]: + return + + cred = self._transport._credentials + + # get_cred_info is only available in google-auth>=2.35.0 + if not hasattr(cred, "get_cred_info"): + return + + # ignore the type check since pypy test fails when get_cred_info + # is not available + cred_info = cred.get_cred_info() # type: ignore + if cred_info and hasattr(error._details, "append"): + error._details.append(json.dumps(cred_info)) + @property def api_endpoint(self): """Return the API endpoint used by the client instance. @@ -2407,16 +2436,20 @@ def list_operations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_operation( self, @@ -2462,16 +2495,20 @@ def get_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def delete_operation( self, @@ -2634,16 +2671,20 @@ def wait_operation( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def set_iam_policy( self, @@ -2755,16 +2796,20 @@ def set_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_iam_policy( self, @@ -2877,16 +2922,20 @@ def get_iam_policy( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def test_iam_permissions( self, @@ -2937,16 +2986,20 @@ def test_iam_permissions( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def get_location( self, @@ -2992,16 +3045,20 @@ def get_location( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e def list_locations( self, @@ -3047,16 +3104,20 @@ def list_locations( # Validate the universe domain. self._validate_universe_domain() - # Send the request. - response = rpc( - request, - retry=retry, - timeout=timeout, - metadata=metadata, - ) + try: + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) - # Done; return the response. - return response + # Done; return the response. + return response + except core_exceptions.GoogleAPICallError as e: + self._add_cred_info_for_auth_errors(e) + raise e DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( diff --git a/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest.py b/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest.py index 47cc6eeeba..ebd60b5048 100644 --- a/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest.py +++ b/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest.py @@ -218,12 +218,33 @@ def pre_add_trial_measurement( def post_add_trial_measurement(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for add_trial_measurement - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_trial_measurement_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_add_trial_measurement` interceptor runs + before the `post_add_trial_measurement_with_metadata` interceptor. """ return response + def post_add_trial_measurement_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for add_trial_measurement + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_add_trial_measurement_with_metadata` + interceptor in new development instead of the `post_add_trial_measurement` interceptor. + When both interceptors are used, this `post_add_trial_measurement_with_metadata` interceptor runs after the + `post_add_trial_measurement` interceptor. The (possibly modified) response returned by + `post_add_trial_measurement` will be passed to + `post_add_trial_measurement_with_metadata`. + """ + return response, metadata + def pre_check_trial_early_stopping_state( self, request: vizier_service.CheckTrialEarlyStoppingStateRequest, @@ -244,12 +265,35 @@ def post_check_trial_early_stopping_state( ) -> operations_pb2.Operation: """Post-rpc interceptor for check_trial_early_stopping_state - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_check_trial_early_stopping_state_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_check_trial_early_stopping_state` interceptor runs + before the `post_check_trial_early_stopping_state_with_metadata` interceptor. """ return response + def post_check_trial_early_stopping_state_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for check_trial_early_stopping_state + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_check_trial_early_stopping_state_with_metadata` + interceptor in new development instead of the `post_check_trial_early_stopping_state` interceptor. + When both interceptors are used, this `post_check_trial_early_stopping_state_with_metadata` interceptor runs after the + `post_check_trial_early_stopping_state` interceptor. The (possibly modified) response returned by + `post_check_trial_early_stopping_state` will be passed to + `post_check_trial_early_stopping_state_with_metadata`. + """ + return response, metadata + def pre_complete_trial( self, request: vizier_service.CompleteTrialRequest, @@ -267,12 +311,33 @@ def pre_complete_trial( def post_complete_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for complete_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_complete_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_complete_trial` interceptor runs + before the `post_complete_trial_with_metadata` interceptor. """ return response + def post_complete_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for complete_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_complete_trial_with_metadata` + interceptor in new development instead of the `post_complete_trial` interceptor. + When both interceptors are used, this `post_complete_trial_with_metadata` interceptor runs after the + `post_complete_trial` interceptor. The (possibly modified) response returned by + `post_complete_trial` will be passed to + `post_complete_trial_with_metadata`. + """ + return response, metadata + def pre_create_study( self, request: vizier_service.CreateStudyRequest, @@ -290,12 +355,35 @@ def pre_create_study( def post_create_study(self, response: gca_study.Study) -> gca_study.Study: """Post-rpc interceptor for create_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_study` interceptor runs + before the `post_create_study_with_metadata` interceptor. """ return response + def post_create_study_with_metadata( + self, + response: gca_study.Study, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_study_with_metadata` + interceptor in new development instead of the `post_create_study` interceptor. + When both interceptors are used, this `post_create_study_with_metadata` interceptor runs after the + `post_create_study` interceptor. The (possibly modified) response returned by + `post_create_study` will be passed to + `post_create_study_with_metadata`. + """ + return response, metadata + def pre_create_trial( self, request: vizier_service.CreateTrialRequest, @@ -313,12 +401,33 @@ def pre_create_trial( def post_create_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for create_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_trial` interceptor runs + before the `post_create_trial_with_metadata` interceptor. """ return response + def post_create_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_trial_with_metadata` + interceptor in new development instead of the `post_create_trial` interceptor. + When both interceptors are used, this `post_create_trial_with_metadata` interceptor runs after the + `post_create_trial` interceptor. The (possibly modified) response returned by + `post_create_trial` will be passed to + `post_create_trial_with_metadata`. + """ + return response, metadata + def pre_delete_study( self, request: vizier_service.DeleteStudyRequest, @@ -362,12 +471,33 @@ def pre_get_study( def post_get_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for get_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_study` interceptor runs + before the `post_get_study_with_metadata` interceptor. """ return response + def post_get_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_study_with_metadata` + interceptor in new development instead of the `post_get_study` interceptor. + When both interceptors are used, this `post_get_study_with_metadata` interceptor runs after the + `post_get_study` interceptor. The (possibly modified) response returned by + `post_get_study` will be passed to + `post_get_study_with_metadata`. + """ + return response, metadata + def pre_get_trial( self, request: vizier_service.GetTrialRequest, @@ -383,12 +513,33 @@ def pre_get_trial( def post_get_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for get_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_trial` interceptor runs + before the `post_get_trial_with_metadata` interceptor. """ return response + def post_get_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_trial_with_metadata` + interceptor in new development instead of the `post_get_trial` interceptor. + When both interceptors are used, this `post_get_trial_with_metadata` interceptor runs after the + `post_get_trial` interceptor. The (possibly modified) response returned by + `post_get_trial` will be passed to + `post_get_trial_with_metadata`. + """ + return response, metadata + def pre_list_optimal_trials( self, request: vizier_service.ListOptimalTrialsRequest, @@ -408,12 +559,38 @@ def post_list_optimal_trials( ) -> vizier_service.ListOptimalTrialsResponse: """Post-rpc interceptor for list_optimal_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_optimal_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_optimal_trials` interceptor runs + before the `post_list_optimal_trials_with_metadata` interceptor. """ return response + def post_list_optimal_trials_with_metadata( + self, + response: vizier_service.ListOptimalTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListOptimalTrialsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_optimal_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_optimal_trials_with_metadata` + interceptor in new development instead of the `post_list_optimal_trials` interceptor. + When both interceptors are used, this `post_list_optimal_trials_with_metadata` interceptor runs after the + `post_list_optimal_trials` interceptor. The (possibly modified) response returned by + `post_list_optimal_trials` will be passed to + `post_list_optimal_trials_with_metadata`. + """ + return response, metadata + def pre_list_studies( self, request: vizier_service.ListStudiesRequest, @@ -433,12 +610,37 @@ def post_list_studies( ) -> vizier_service.ListStudiesResponse: """Post-rpc interceptor for list_studies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_studies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_studies` interceptor runs + before the `post_list_studies_with_metadata` interceptor. """ return response + def post_list_studies_with_metadata( + self, + response: vizier_service.ListStudiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListStudiesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_studies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_studies_with_metadata` + interceptor in new development instead of the `post_list_studies` interceptor. + When both interceptors are used, this `post_list_studies_with_metadata` interceptor runs after the + `post_list_studies` interceptor. The (possibly modified) response returned by + `post_list_studies` will be passed to + `post_list_studies_with_metadata`. + """ + return response, metadata + def pre_list_trials( self, request: vizier_service.ListTrialsRequest, @@ -458,12 +660,37 @@ def post_list_trials( ) -> vizier_service.ListTrialsResponse: """Post-rpc interceptor for list_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_trials` interceptor runs + before the `post_list_trials_with_metadata` interceptor. """ return response + def post_list_trials_with_metadata( + self, + response: vizier_service.ListTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListTrialsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_trials_with_metadata` + interceptor in new development instead of the `post_list_trials` interceptor. + When both interceptors are used, this `post_list_trials_with_metadata` interceptor runs after the + `post_list_trials` interceptor. The (possibly modified) response returned by + `post_list_trials` will be passed to + `post_list_trials_with_metadata`. + """ + return response, metadata + def pre_lookup_study( self, request: vizier_service.LookupStudyRequest, @@ -481,12 +708,33 @@ def pre_lookup_study( def post_lookup_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for lookup_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_lookup_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_lookup_study` interceptor runs + before the `post_lookup_study_with_metadata` interceptor. """ return response + def post_lookup_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for lookup_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_lookup_study_with_metadata` + interceptor in new development instead of the `post_lookup_study` interceptor. + When both interceptors are used, this `post_lookup_study_with_metadata` interceptor runs after the + `post_lookup_study` interceptor. The (possibly modified) response returned by + `post_lookup_study` will be passed to + `post_lookup_study_with_metadata`. + """ + return response, metadata + def pre_stop_trial( self, request: vizier_service.StopTrialRequest, @@ -504,12 +752,33 @@ def pre_stop_trial( def post_stop_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for stop_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_stop_trial` interceptor runs + before the `post_stop_trial_with_metadata` interceptor. """ return response + def post_stop_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_stop_trial_with_metadata` + interceptor in new development instead of the `post_stop_trial` interceptor. + When both interceptors are used, this `post_stop_trial_with_metadata` interceptor runs after the + `post_stop_trial` interceptor. The (possibly modified) response returned by + `post_stop_trial` will be passed to + `post_stop_trial_with_metadata`. + """ + return response, metadata + def pre_suggest_trials( self, request: vizier_service.SuggestTrialsRequest, @@ -529,12 +798,35 @@ def post_suggest_trials( ) -> operations_pb2.Operation: """Post-rpc interceptor for suggest_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_suggest_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_suggest_trials` interceptor runs + before the `post_suggest_trials_with_metadata` interceptor. """ return response + def post_suggest_trials_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for suggest_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_suggest_trials_with_metadata` + interceptor in new development instead of the `post_suggest_trials` interceptor. + When both interceptors are used, this `post_suggest_trials_with_metadata` interceptor runs after the + `post_suggest_trials` interceptor. The (possibly modified) response returned by + `post_suggest_trials` will be passed to + `post_suggest_trials_with_metadata`. + """ + return response, metadata + def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -3088,6 +3380,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_add_trial_measurement(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_add_trial_measurement_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3243,6 +3539,13 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_check_trial_early_stopping_state(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = self._interceptor.post_check_trial_early_stopping_state_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3396,6 +3699,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_complete_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_complete_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3548,6 +3855,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3705,6 +4016,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_create_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_create_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4071,6 +4386,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4222,6 +4541,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_get_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_get_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4374,6 +4697,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_optimal_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_optimal_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4522,6 +4849,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_studies(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_studies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4672,6 +5003,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_list_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_list_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4826,6 +5161,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_lookup_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_lookup_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -4985,6 +5324,10 @@ def __call__( json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) resp = self._interceptor.post_stop_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_stop_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -5134,6 +5477,10 @@ def __call__( json_format.Parse(response.content, resp, ignore_unknown_fields=True) resp = self._interceptor.post_suggest_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = self._interceptor.post_suggest_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest_asyncio.py b/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest_asyncio.py index a472258357..20fb1be783 100644 --- a/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest_asyncio.py +++ b/google/cloud/aiplatform_v1beta1/services/vizier_service/transports/rest_asyncio.py @@ -236,12 +236,33 @@ async def pre_add_trial_measurement( async def post_add_trial_measurement(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for add_trial_measurement - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_add_trial_measurement_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_add_trial_measurement` interceptor runs + before the `post_add_trial_measurement_with_metadata` interceptor. """ return response + async def post_add_trial_measurement_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for add_trial_measurement + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_add_trial_measurement_with_metadata` + interceptor in new development instead of the `post_add_trial_measurement` interceptor. + When both interceptors are used, this `post_add_trial_measurement_with_metadata` interceptor runs after the + `post_add_trial_measurement` interceptor. The (possibly modified) response returned by + `post_add_trial_measurement` will be passed to + `post_add_trial_measurement_with_metadata`. + """ + return response, metadata + async def pre_check_trial_early_stopping_state( self, request: vizier_service.CheckTrialEarlyStoppingStateRequest, @@ -262,12 +283,35 @@ async def post_check_trial_early_stopping_state( ) -> operations_pb2.Operation: """Post-rpc interceptor for check_trial_early_stopping_state - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_check_trial_early_stopping_state_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_check_trial_early_stopping_state` interceptor runs + before the `post_check_trial_early_stopping_state_with_metadata` interceptor. """ return response + async def post_check_trial_early_stopping_state_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for check_trial_early_stopping_state + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_check_trial_early_stopping_state_with_metadata` + interceptor in new development instead of the `post_check_trial_early_stopping_state` interceptor. + When both interceptors are used, this `post_check_trial_early_stopping_state_with_metadata` interceptor runs after the + `post_check_trial_early_stopping_state` interceptor. The (possibly modified) response returned by + `post_check_trial_early_stopping_state` will be passed to + `post_check_trial_early_stopping_state_with_metadata`. + """ + return response, metadata + async def pre_complete_trial( self, request: vizier_service.CompleteTrialRequest, @@ -285,12 +329,33 @@ async def pre_complete_trial( async def post_complete_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for complete_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_complete_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_complete_trial` interceptor runs + before the `post_complete_trial_with_metadata` interceptor. """ return response + async def post_complete_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for complete_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_complete_trial_with_metadata` + interceptor in new development instead of the `post_complete_trial` interceptor. + When both interceptors are used, this `post_complete_trial_with_metadata` interceptor runs after the + `post_complete_trial` interceptor. The (possibly modified) response returned by + `post_complete_trial` will be passed to + `post_complete_trial_with_metadata`. + """ + return response, metadata + async def pre_create_study( self, request: vizier_service.CreateStudyRequest, @@ -308,12 +373,35 @@ async def pre_create_study( async def post_create_study(self, response: gca_study.Study) -> gca_study.Study: """Post-rpc interceptor for create_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_study` interceptor runs + before the `post_create_study_with_metadata` interceptor. """ return response + async def post_create_study_with_metadata( + self, + response: gca_study.Study, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[gca_study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_study_with_metadata` + interceptor in new development instead of the `post_create_study` interceptor. + When both interceptors are used, this `post_create_study_with_metadata` interceptor runs after the + `post_create_study` interceptor. The (possibly modified) response returned by + `post_create_study` will be passed to + `post_create_study_with_metadata`. + """ + return response, metadata + async def pre_create_trial( self, request: vizier_service.CreateTrialRequest, @@ -331,12 +419,33 @@ async def pre_create_trial( async def post_create_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for create_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_create_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_create_trial` interceptor runs + before the `post_create_trial_with_metadata` interceptor. """ return response + async def post_create_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for create_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_create_trial_with_metadata` + interceptor in new development instead of the `post_create_trial` interceptor. + When both interceptors are used, this `post_create_trial_with_metadata` interceptor runs after the + `post_create_trial` interceptor. The (possibly modified) response returned by + `post_create_trial` will be passed to + `post_create_trial_with_metadata`. + """ + return response, metadata + async def pre_delete_study( self, request: vizier_service.DeleteStudyRequest, @@ -380,12 +489,33 @@ async def pre_get_study( async def post_get_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for get_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_study` interceptor runs + before the `post_get_study_with_metadata` interceptor. """ return response + async def post_get_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_study_with_metadata` + interceptor in new development instead of the `post_get_study` interceptor. + When both interceptors are used, this `post_get_study_with_metadata` interceptor runs after the + `post_get_study` interceptor. The (possibly modified) response returned by + `post_get_study` will be passed to + `post_get_study_with_metadata`. + """ + return response, metadata + async def pre_get_trial( self, request: vizier_service.GetTrialRequest, @@ -401,12 +531,33 @@ async def pre_get_trial( async def post_get_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for get_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_get_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_get_trial` interceptor runs + before the `post_get_trial_with_metadata` interceptor. """ return response + async def post_get_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for get_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_get_trial_with_metadata` + interceptor in new development instead of the `post_get_trial` interceptor. + When both interceptors are used, this `post_get_trial_with_metadata` interceptor runs after the + `post_get_trial` interceptor. The (possibly modified) response returned by + `post_get_trial` will be passed to + `post_get_trial_with_metadata`. + """ + return response, metadata + async def pre_list_optimal_trials( self, request: vizier_service.ListOptimalTrialsRequest, @@ -426,12 +577,38 @@ async def post_list_optimal_trials( ) -> vizier_service.ListOptimalTrialsResponse: """Post-rpc interceptor for list_optimal_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_optimal_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_optimal_trials` interceptor runs + before the `post_list_optimal_trials_with_metadata` interceptor. """ return response + async def post_list_optimal_trials_with_metadata( + self, + response: vizier_service.ListOptimalTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListOptimalTrialsResponse, + Sequence[Tuple[str, Union[str, bytes]]], + ]: + """Post-rpc interceptor for list_optimal_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_optimal_trials_with_metadata` + interceptor in new development instead of the `post_list_optimal_trials` interceptor. + When both interceptors are used, this `post_list_optimal_trials_with_metadata` interceptor runs after the + `post_list_optimal_trials` interceptor. The (possibly modified) response returned by + `post_list_optimal_trials` will be passed to + `post_list_optimal_trials_with_metadata`. + """ + return response, metadata + async def pre_list_studies( self, request: vizier_service.ListStudiesRequest, @@ -451,12 +628,37 @@ async def post_list_studies( ) -> vizier_service.ListStudiesResponse: """Post-rpc interceptor for list_studies - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_studies_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_studies` interceptor runs + before the `post_list_studies_with_metadata` interceptor. """ return response + async def post_list_studies_with_metadata( + self, + response: vizier_service.ListStudiesResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListStudiesResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_studies + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_studies_with_metadata` + interceptor in new development instead of the `post_list_studies` interceptor. + When both interceptors are used, this `post_list_studies_with_metadata` interceptor runs after the + `post_list_studies` interceptor. The (possibly modified) response returned by + `post_list_studies` will be passed to + `post_list_studies_with_metadata`. + """ + return response, metadata + async def pre_list_trials( self, request: vizier_service.ListTrialsRequest, @@ -476,12 +678,37 @@ async def post_list_trials( ) -> vizier_service.ListTrialsResponse: """Post-rpc interceptor for list_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_list_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_list_trials` interceptor runs + before the `post_list_trials_with_metadata` interceptor. """ return response + async def post_list_trials_with_metadata( + self, + response: vizier_service.ListTrialsResponse, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[ + vizier_service.ListTrialsResponse, Sequence[Tuple[str, Union[str, bytes]]] + ]: + """Post-rpc interceptor for list_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_list_trials_with_metadata` + interceptor in new development instead of the `post_list_trials` interceptor. + When both interceptors are used, this `post_list_trials_with_metadata` interceptor runs after the + `post_list_trials` interceptor. The (possibly modified) response returned by + `post_list_trials` will be passed to + `post_list_trials_with_metadata`. + """ + return response, metadata + async def pre_lookup_study( self, request: vizier_service.LookupStudyRequest, @@ -499,12 +726,33 @@ async def pre_lookup_study( async def post_lookup_study(self, response: study.Study) -> study.Study: """Post-rpc interceptor for lookup_study - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_lookup_study_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_lookup_study` interceptor runs + before the `post_lookup_study_with_metadata` interceptor. """ return response + async def post_lookup_study_with_metadata( + self, response: study.Study, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Study, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for lookup_study + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_lookup_study_with_metadata` + interceptor in new development instead of the `post_lookup_study` interceptor. + When both interceptors are used, this `post_lookup_study_with_metadata` interceptor runs after the + `post_lookup_study` interceptor. The (possibly modified) response returned by + `post_lookup_study` will be passed to + `post_lookup_study_with_metadata`. + """ + return response, metadata + async def pre_stop_trial( self, request: vizier_service.StopTrialRequest, @@ -522,12 +770,33 @@ async def pre_stop_trial( async def post_stop_trial(self, response: study.Trial) -> study.Trial: """Post-rpc interceptor for stop_trial - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_stop_trial_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_stop_trial` interceptor runs + before the `post_stop_trial_with_metadata` interceptor. """ return response + async def post_stop_trial_with_metadata( + self, response: study.Trial, metadata: Sequence[Tuple[str, Union[str, bytes]]] + ) -> Tuple[study.Trial, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for stop_trial + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_stop_trial_with_metadata` + interceptor in new development instead of the `post_stop_trial` interceptor. + When both interceptors are used, this `post_stop_trial_with_metadata` interceptor runs after the + `post_stop_trial` interceptor. The (possibly modified) response returned by + `post_stop_trial` will be passed to + `post_stop_trial_with_metadata`. + """ + return response, metadata + async def pre_suggest_trials( self, request: vizier_service.SuggestTrialsRequest, @@ -547,12 +816,35 @@ async def post_suggest_trials( ) -> operations_pb2.Operation: """Post-rpc interceptor for suggest_trials - Override in a subclass to manipulate the response + DEPRECATED. Please use the `post_suggest_trials_with_metadata` + interceptor instead. + + Override in a subclass to read or manipulate the response after it is returned by the VizierService server but before - it is returned to user code. + it is returned to user code. This `post_suggest_trials` interceptor runs + before the `post_suggest_trials_with_metadata` interceptor. """ return response + async def post_suggest_trials_with_metadata( + self, + response: operations_pb2.Operation, + metadata: Sequence[Tuple[str, Union[str, bytes]]], + ) -> Tuple[operations_pb2.Operation, Sequence[Tuple[str, Union[str, bytes]]]]: + """Post-rpc interceptor for suggest_trials + + Override in a subclass to read or manipulate the response or metadata after it + is returned by the VizierService server but before it is returned to user code. + + We recommend only using this `post_suggest_trials_with_metadata` + interceptor in new development instead of the `post_suggest_trials` interceptor. + When both interceptors are used, this `post_suggest_trials_with_metadata` interceptor runs after the + `post_suggest_trials` interceptor. The (possibly modified) response returned by + `post_suggest_trials` will be passed to + `post_suggest_trials_with_metadata`. + """ + return response, metadata + async def pre_get_location( self, request: locations_pb2.GetLocationRequest, @@ -1142,6 +1434,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_add_trial_measurement(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_add_trial_measurement_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1306,6 +1602,13 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_check_trial_early_stopping_state(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + ( + resp, + _, + ) = await self._interceptor.post_check_trial_early_stopping_state_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1469,6 +1772,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_complete_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_complete_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1629,6 +1936,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -1794,6 +2105,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_create_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_create_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2182,6 +2497,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2339,6 +2658,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_get_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_get_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2500,6 +2823,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_optimal_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_optimal_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2656,6 +2983,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_studies(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_studies_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2814,6 +3145,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_list_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_list_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -2976,6 +3311,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_lookup_study(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_lookup_study_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3143,6 +3482,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_stop_trial(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_stop_trial_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER @@ -3304,6 +3647,10 @@ async def __call__( content = await response.read() json_format.Parse(content, pb_resp, ignore_unknown_fields=True) resp = await self._interceptor.post_suggest_trials(resp) + response_metadata = [(k, str(v)) for k, v in response.headers.items()] + resp, _ = await self._interceptor.post_suggest_trials_with_metadata( + resp, response_metadata + ) if CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor( logging.DEBUG ): # pragma: NO COVER diff --git a/google/cloud/aiplatform_v1beta1/types/__init__.py b/google/cloud/aiplatform_v1beta1/types/__init__.py index 4a264c643d..042f984c25 100644 --- a/google/cloud/aiplatform_v1beta1/types/__init__.py +++ b/google/cloud/aiplatform_v1beta1/types/__init__.py @@ -158,6 +158,7 @@ FasterDeploymentConfig, PredictRequestResponseLoggingConfig, PrivateEndpoints, + RolloutOptions, ) from .endpoint_service import ( CreateEndpointOperationMetadata, @@ -205,8 +206,12 @@ CometInstance, CometResult, CometSpec, + EvaluateDatasetOperationMetadata, + EvaluateDatasetRequest, + EvaluateDatasetResponse, EvaluateInstancesRequest, EvaluateInstancesResponse, + EvaluationDataset, ExactMatchInput, ExactMatchInstance, ExactMatchMetricValue, @@ -224,10 +229,13 @@ GroundednessInstance, GroundednessResult, GroundednessSpec, + Metric, MetricxInput, MetricxInstance, MetricxResult, MetricxSpec, + OutputConfig, + OutputInfo, PairwiseMetricInput, PairwiseMetricInstance, PairwiseMetricResult, @@ -878,9 +886,12 @@ ListModelEvaluationsResponse, ListModelsRequest, ListModelsResponse, + ListModelVersionCheckpointsRequest, + ListModelVersionCheckpointsResponse, ListModelVersionsRequest, ListModelVersionsResponse, MergeVersionAliasesRequest, + ModelVersionCheckpoint, UpdateExplanationDatasetOperationMetadata, UpdateExplanationDatasetRequest, UpdateExplanationDatasetResponse, @@ -947,6 +958,10 @@ UpgradeNotebookRuntimeResponse, NotebookExecutionJobView, ) +from .notebook_software_config import ( + NotebookSoftwareConfig, + PostStartupScriptConfig, +) from .openapi import ( Schema, Type, @@ -1410,6 +1425,7 @@ "FasterDeploymentConfig", "PredictRequestResponseLoggingConfig", "PrivateEndpoints", + "RolloutOptions", "CreateEndpointOperationMetadata", "CreateEndpointRequest", "DeleteEndpointRequest", @@ -1447,8 +1463,12 @@ "CometInstance", "CometResult", "CometSpec", + "EvaluateDatasetOperationMetadata", + "EvaluateDatasetRequest", + "EvaluateDatasetResponse", "EvaluateInstancesRequest", "EvaluateInstancesResponse", + "EvaluationDataset", "ExactMatchInput", "ExactMatchInstance", "ExactMatchMetricValue", @@ -1466,10 +1486,13 @@ "GroundednessInstance", "GroundednessResult", "GroundednessSpec", + "Metric", "MetricxInput", "MetricxInstance", "MetricxResult", "MetricxSpec", + "OutputConfig", + "OutputInfo", "PairwiseMetricInput", "PairwiseMetricInstance", "PairwiseMetricResult", @@ -2008,9 +2031,12 @@ "ListModelEvaluationsResponse", "ListModelsRequest", "ListModelsResponse", + "ListModelVersionCheckpointsRequest", + "ListModelVersionCheckpointsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", "MergeVersionAliasesRequest", + "ModelVersionCheckpoint", "UpdateExplanationDatasetOperationMetadata", "UpdateExplanationDatasetRequest", "UpdateExplanationDatasetResponse", @@ -2060,6 +2086,8 @@ "UpgradeNotebookRuntimeRequest", "UpgradeNotebookRuntimeResponse", "NotebookExecutionJobView", + "NotebookSoftwareConfig", + "PostStartupScriptConfig", "Schema", "Type", "DeleteOperationMetadata", diff --git a/google/cloud/aiplatform_v1beta1/types/endpoint.py b/google/cloud/aiplatform_v1beta1/types/endpoint.py index 0672f8d0f0..009d2d39f4 100644 --- a/google/cloud/aiplatform_v1beta1/types/endpoint.py +++ b/google/cloud/aiplatform_v1beta1/types/endpoint.py @@ -37,6 +37,7 @@ "PredictRequestResponseLoggingConfig", "ClientConnectionConfig", "FasterDeploymentConfig", + "RolloutOptions", }, ) @@ -377,6 +378,8 @@ class DeployedModel(proto.Message): is configured. faster_deployment_config (google.cloud.aiplatform_v1beta1.types.FasterDeploymentConfig): Configuration for faster model deployment. + rollout_options (google.cloud.aiplatform_v1beta1.types.RolloutOptions): + Options for configuring rolling deployments. status (google.cloud.aiplatform_v1beta1.types.DeployedModel.Status): Output only. Runtime status of the deployed model. @@ -484,6 +487,11 @@ class Status(proto.Message): number=23, message="FasterDeploymentConfig", ) + rollout_options: "RolloutOptions" = proto.Field( + proto.MESSAGE, + number=25, + message="RolloutOptions", + ) status: Status = proto.Field( proto.MESSAGE, number=26, @@ -603,4 +611,78 @@ class FasterDeploymentConfig(proto.Message): ) +class RolloutOptions(proto.Message): + r"""Configuration for rolling deployments. + + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + max_unavailable_replicas (int): + Absolute count of replicas allowed to be + unavailable. + + This field is a member of `oneof`_ ``max_unavailable``. + max_unavailable_percentage (int): + Percentage of replicas allowed to be + unavailable. For autoscaling deployments, this + refers to the target replica count. + + This field is a member of `oneof`_ ``max_unavailable``. + max_surge_replicas (int): + Absolute count of allowed additional + replicas. + + This field is a member of `oneof`_ ``max_surge``. + max_surge_percentage (int): + Percentage of allowed additional replicas. + For autoscaling deployments, this refers to the + target replica count. + + This field is a member of `oneof`_ ``max_surge``. + previous_deployed_model (str): + ID of the DeployedModel that this deployment + should replace. + revision_number (int): + Output only. Read-only. Revision number + determines the relative priority of + DeployedModels in the same rollout. The + DeployedModel with the largest revision number + specifies the intended state of the deployment. + """ + + max_unavailable_replicas: int = proto.Field( + proto.INT32, + number=3, + oneof="max_unavailable", + ) + max_unavailable_percentage: int = proto.Field( + proto.INT32, + number=4, + oneof="max_unavailable", + ) + max_surge_replicas: int = proto.Field( + proto.INT32, + number=5, + oneof="max_surge", + ) + max_surge_percentage: int = proto.Field( + proto.INT32, + number=6, + oneof="max_surge", + ) + previous_deployed_model: str = proto.Field( + proto.STRING, + number=1, + ) + revision_number: int = proto.Field( + proto.INT32, + number=2, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/aiplatform_v1beta1/types/evaluation_service.py b/google/cloud/aiplatform_v1beta1/types/evaluation_service.py index c4109d5943..912b0d6d92 100644 --- a/google/cloud/aiplatform_v1beta1/types/evaluation_service.py +++ b/google/cloud/aiplatform_v1beta1/types/evaluation_service.py @@ -19,11 +19,21 @@ import proto # type: ignore +from google.cloud.aiplatform_v1beta1.types import io +from google.cloud.aiplatform_v1beta1.types import operation + __protobuf__ = proto.module( package="google.cloud.aiplatform.v1beta1", manifest={ "PairwiseChoice", + "EvaluateDatasetOperationMetadata", + "EvaluateDatasetResponse", + "OutputInfo", + "EvaluateDatasetRequest", + "OutputConfig", + "Metric", + "EvaluationDataset", "AutoraterConfig", "EvaluateInstancesRequest", "EvaluateInstancesResponse", @@ -189,6 +199,273 @@ class PairwiseChoice(proto.Enum): TIE = 3 +class EvaluateDatasetOperationMetadata(proto.Message): + r"""Operation metadata for Dataset Evaluation. + + Attributes: + generic_metadata (google.cloud.aiplatform_v1beta1.types.GenericOperationMetadata): + Generic operation metadata. + """ + + generic_metadata: operation.GenericOperationMetadata = proto.Field( + proto.MESSAGE, + number=1, + message=operation.GenericOperationMetadata, + ) + + +class EvaluateDatasetResponse(proto.Message): + r"""Response in LRO for EvaluationService.EvaluateDataset. + + Attributes: + output_info (google.cloud.aiplatform_v1beta1.types.OutputInfo): + Output only. Output info for + EvaluationService.EvaluateDataset. + """ + + output_info: "OutputInfo" = proto.Field( + proto.MESSAGE, + number=3, + message="OutputInfo", + ) + + +class OutputInfo(proto.Message): + r"""Describes the info for output of + EvaluationService.EvaluateDataset. + + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + gcs_output_directory (str): + Output only. The full path of the Cloud + Storage directory created, into which the + evaluation results and aggregation results are + written. + + This field is a member of `oneof`_ ``output_location``. + """ + + gcs_output_directory: str = proto.Field( + proto.STRING, + number=1, + oneof="output_location", + ) + + +class EvaluateDatasetRequest(proto.Message): + r"""Request message for EvaluationService.EvaluateDataset. + + Attributes: + location (str): + Required. The resource name of the Location to evaluate the + dataset. Format: ``projects/{project}/locations/{location}`` + dataset (google.cloud.aiplatform_v1beta1.types.EvaluationDataset): + Required. The dataset used for evaluation. + metrics (MutableSequence[google.cloud.aiplatform_v1beta1.types.Metric]): + Required. The metrics used for evaluation. + output_config (google.cloud.aiplatform_v1beta1.types.OutputConfig): + Required. Config for evaluation output. + autorater_config (google.cloud.aiplatform_v1beta1.types.AutoraterConfig): + Optional. Autorater config used for + evaluation. + """ + + location: str = proto.Field( + proto.STRING, + number=1, + ) + dataset: "EvaluationDataset" = proto.Field( + proto.MESSAGE, + number=2, + message="EvaluationDataset", + ) + metrics: MutableSequence["Metric"] = proto.RepeatedField( + proto.MESSAGE, + number=3, + message="Metric", + ) + output_config: "OutputConfig" = proto.Field( + proto.MESSAGE, + number=4, + message="OutputConfig", + ) + autorater_config: "AutoraterConfig" = proto.Field( + proto.MESSAGE, + number=5, + message="AutoraterConfig", + ) + + +class OutputConfig(proto.Message): + r"""Config for evaluation output. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + gcs_destination (google.cloud.aiplatform_v1beta1.types.GcsDestination): + Cloud storage destination for evaluation + output. + + This field is a member of `oneof`_ ``destination``. + """ + + gcs_destination: io.GcsDestination = proto.Field( + proto.MESSAGE, + number=1, + oneof="destination", + message=io.GcsDestination, + ) + + +class Metric(proto.Message): + r"""The metric used for dataset level evaluation. + + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + pointwise_metric_spec (google.cloud.aiplatform_v1beta1.types.PointwiseMetricSpec): + Spec for pointwise metric. + + This field is a member of `oneof`_ ``metric_spec``. + pairwise_metric_spec (google.cloud.aiplatform_v1beta1.types.PairwiseMetricSpec): + Spec for pairwise metric. + + This field is a member of `oneof`_ ``metric_spec``. + exact_match_spec (google.cloud.aiplatform_v1beta1.types.ExactMatchSpec): + Spec for exact match metric. + + This field is a member of `oneof`_ ``metric_spec``. + bleu_spec (google.cloud.aiplatform_v1beta1.types.BleuSpec): + Spec for bleu metric. + + This field is a member of `oneof`_ ``metric_spec``. + rouge_spec (google.cloud.aiplatform_v1beta1.types.RougeSpec): + Spec for rouge metric. + + This field is a member of `oneof`_ ``metric_spec``. + aggregation_metrics (MutableSequence[google.cloud.aiplatform_v1beta1.types.Metric.AggregationMetric]): + Optional. The aggregation metrics to use. + """ + + class AggregationMetric(proto.Enum): + r"""The aggregation metrics supported by + EvaluationService.EvaluateDataset. + + Values: + AGGREGATION_METRIC_UNSPECIFIED (0): + Unspecified aggregation metric. + AVERAGE (1): + Average aggregation metric. + MODE (2): + Mode aggregation metric. + STANDARD_DEVIATION (3): + Standard deviation aggregation metric. + VARIANCE (4): + Variance aggregation metric. + MINIMUM (5): + Minimum aggregation metric. + MAXIMUM (6): + Maximum aggregation metric. + MEDIAN (7): + Median aggregation metric. + PERCENTILE_P90 (8): + 90th percentile aggregation metric. + PERCENTILE_P95 (9): + 95th percentile aggregation metric. + PERCENTILE_P99 (10): + 99th percentile aggregation metric. + """ + AGGREGATION_METRIC_UNSPECIFIED = 0 + AVERAGE = 1 + MODE = 2 + STANDARD_DEVIATION = 3 + VARIANCE = 4 + MINIMUM = 5 + MAXIMUM = 6 + MEDIAN = 7 + PERCENTILE_P90 = 8 + PERCENTILE_P95 = 9 + PERCENTILE_P99 = 10 + + pointwise_metric_spec: "PointwiseMetricSpec" = proto.Field( + proto.MESSAGE, + number=2, + oneof="metric_spec", + message="PointwiseMetricSpec", + ) + pairwise_metric_spec: "PairwiseMetricSpec" = proto.Field( + proto.MESSAGE, + number=3, + oneof="metric_spec", + message="PairwiseMetricSpec", + ) + exact_match_spec: "ExactMatchSpec" = proto.Field( + proto.MESSAGE, + number=4, + oneof="metric_spec", + message="ExactMatchSpec", + ) + bleu_spec: "BleuSpec" = proto.Field( + proto.MESSAGE, + number=5, + oneof="metric_spec", + message="BleuSpec", + ) + rouge_spec: "RougeSpec" = proto.Field( + proto.MESSAGE, + number=6, + oneof="metric_spec", + message="RougeSpec", + ) + aggregation_metrics: MutableSequence[AggregationMetric] = proto.RepeatedField( + proto.ENUM, + number=1, + enum=AggregationMetric, + ) + + +class EvaluationDataset(proto.Message): + r"""The dataset used for evaluation. + + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + gcs_source (google.cloud.aiplatform_v1beta1.types.GcsSource): + Cloud storage source holds the dataset. + + This field is a member of `oneof`_ ``source``. + bigquery_source (google.cloud.aiplatform_v1beta1.types.BigQuerySource): + BigQuery source holds the dataset. + + This field is a member of `oneof`_ ``source``. + """ + + gcs_source: io.GcsSource = proto.Field( + proto.MESSAGE, + number=1, + oneof="source", + message=io.GcsSource, + ) + bigquery_source: io.BigQuerySource = proto.Field( + proto.MESSAGE, + number=2, + oneof="source", + message=io.BigQuerySource, + ) + + class AutoraterConfig(proto.Message): r"""The configs for autorater. This is applicable to both EvaluateInstances and EvaluateDataset. diff --git a/google/cloud/aiplatform_v1beta1/types/model.py b/google/cloud/aiplatform_v1beta1/types/model.py index 3db12ad1fb..2dea38b9d2 100644 --- a/google/cloud/aiplatform_v1beta1/types/model.py +++ b/google/cloud/aiplatform_v1beta1/types/model.py @@ -80,6 +80,8 @@ class Model(proto.Message): The description of the Model. version_description (str): The description of this version. + default_checkpoint_id (str): + The default checkpoint id of a model version. predict_schemata (google.cloud.aiplatform_v1beta1.types.PredictSchemata): The schemata that describe formats of the Model's predictions and explanations as given and returned via @@ -492,6 +494,10 @@ class BaseModelSource(proto.Message): proto.STRING, number=30, ) + default_checkpoint_id: str = proto.Field( + proto.STRING, + number=53, + ) predict_schemata: "PredictSchemata" = proto.Field( proto.MESSAGE, number=4, @@ -983,6 +989,9 @@ class ModelContainerSpec(proto.Message): health_probe (google.cloud.aiplatform_v1beta1.types.Probe): Immutable. Specification for Kubernetes readiness probe. + liveness_probe (google.cloud.aiplatform_v1beta1.types.Probe): + Immutable. Specification for Kubernetes + liveness probe. """ image_uri: str = proto.Field( @@ -1039,6 +1048,11 @@ class ModelContainerSpec(proto.Message): number=13, message="Probe", ) + liveness_probe: "Probe" = proto.Field( + proto.MESSAGE, + number=14, + message="Probe", + ) class Port(proto.Message): @@ -1123,6 +1137,10 @@ class Probe(proto.Message): container to determine whether it is alive or ready to receive traffic. + This message has `oneof`_ fields (mutually exclusive fields). + For each oneof, at most one member field can be set at the same time. + Setting any member of the oneof automatically clears all other + members. .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields @@ -1131,6 +1149,21 @@ class Probe(proto.Message): ExecAction probes the health of a container by executing a command. + This field is a member of `oneof`_ ``probe_type``. + http_get (google.cloud.aiplatform_v1beta1.types.Probe.HttpGetAction): + HttpGetAction probes the health of a + container by sending an HTTP GET request. + + This field is a member of `oneof`_ ``probe_type``. + grpc (google.cloud.aiplatform_v1beta1.types.Probe.GrpcAction): + GrpcAction probes the health of a container + by sending a gRPC request. + + This field is a member of `oneof`_ ``probe_type``. + tcp_socket (google.cloud.aiplatform_v1beta1.types.Probe.TcpSocketAction): + TcpSocketAction probes the health of a + container by opening a TCP socket connection. + This field is a member of `oneof`_ ``probe_type``. period_seconds (int): How often (in seconds) to perform the probe. Default to 10 @@ -1144,6 +1177,26 @@ class Probe(proto.Message): period_seconds. Maps to Kubernetes probe argument 'timeoutSeconds'. + failure_threshold (int): + Number of consecutive failures before the + probe is considered failed. Defaults to 3. + Minimum value is 1. + + Maps to Kubernetes probe argument + 'failureThreshold'. + success_threshold (int): + Number of consecutive successes before the + probe is considered successful. Defaults to 1. + Minimum value is 1. + + Maps to Kubernetes probe argument + 'successThreshold'. + initial_delay_seconds (int): + Number of seconds to wait before starting the + probe. Defaults to 0. Minimum value is 0. + + Maps to Kubernetes probe argument + 'initialDelaySeconds'. """ class ExecAction(proto.Message): @@ -1165,12 +1218,147 @@ class ExecAction(proto.Message): number=1, ) + class HttpGetAction(proto.Message): + r"""HttpGetAction describes an action based on HTTP Get requests. + + Attributes: + path (str): + Path to access on the HTTP server. + port (int): + Number of the port to access on the + container. Number must be in the range 1 to + 65535. + host (str): + Host name to connect to, defaults to the + model serving container's IP. You probably want + to set "Host" in httpHeaders instead. + scheme (str): + Scheme to use for connecting to the host. + Defaults to HTTP. Acceptable values are "HTTP" + or "HTTPS". + http_headers (MutableSequence[google.cloud.aiplatform_v1beta1.types.Probe.HttpHeader]): + Custom headers to set in the request. HTTP + allows repeated headers. + """ + + path: str = proto.Field( + proto.STRING, + number=1, + ) + port: int = proto.Field( + proto.INT32, + number=2, + ) + host: str = proto.Field( + proto.STRING, + number=3, + ) + scheme: str = proto.Field( + proto.STRING, + number=4, + ) + http_headers: MutableSequence["Probe.HttpHeader"] = proto.RepeatedField( + proto.MESSAGE, + number=5, + message="Probe.HttpHeader", + ) + + class GrpcAction(proto.Message): + r"""GrpcAction checks the health of a container using a gRPC + service. + + Attributes: + port (int): + Port number of the gRPC service. Number must + be in the range 1 to 65535. + service (str): + Service is the name of the service to place + in the gRPC HealthCheckRequest (see + https://github.com/grpc/grpc/blob/master/doc/health-checking.md). + + If this is not specified, the default behavior + is defined by gRPC. + """ + + port: int = proto.Field( + proto.INT32, + number=1, + ) + service: str = proto.Field( + proto.STRING, + number=2, + ) + + class TcpSocketAction(proto.Message): + r"""TcpSocketAction probes the health of a container by opening a + TCP socket connection. + + Attributes: + port (int): + Number of the port to access on the + container. Number must be in the range 1 to + 65535. + host (str): + Optional: Host name to connect to, defaults + to the model serving container's IP. + """ + + port: int = proto.Field( + proto.INT32, + number=1, + ) + host: str = proto.Field( + proto.STRING, + number=2, + ) + + class HttpHeader(proto.Message): + r"""HttpHeader describes a custom header to be used in HTTP + probes + + Attributes: + name (str): + The header field name. + This will be canonicalized upon output, so + case-variant names will be understood as the + same header. + value (str): + The header field value + """ + + name: str = proto.Field( + proto.STRING, + number=1, + ) + value: str = proto.Field( + proto.STRING, + number=2, + ) + exec_: ExecAction = proto.Field( proto.MESSAGE, number=1, oneof="probe_type", message=ExecAction, ) + http_get: HttpGetAction = proto.Field( + proto.MESSAGE, + number=4, + oneof="probe_type", + message=HttpGetAction, + ) + grpc: GrpcAction = proto.Field( + proto.MESSAGE, + number=5, + oneof="probe_type", + message=GrpcAction, + ) + tcp_socket: TcpSocketAction = proto.Field( + proto.MESSAGE, + number=6, + oneof="probe_type", + message=TcpSocketAction, + ) period_seconds: int = proto.Field( proto.INT32, number=2, @@ -1179,6 +1367,18 @@ class ExecAction(proto.Message): proto.INT32, number=3, ) + failure_threshold: int = proto.Field( + proto.INT32, + number=7, + ) + success_threshold: int = proto.Field( + proto.INT32, + number=8, + ) + initial_delay_seconds: int = proto.Field( + proto.INT32, + number=9, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/aiplatform_v1beta1/types/model_service.py b/google/cloud/aiplatform_v1beta1/types/model_service.py index ffd82f0234..142fffe835 100644 --- a/google/cloud/aiplatform_v1beta1/types/model_service.py +++ b/google/cloud/aiplatform_v1beta1/types/model_service.py @@ -43,6 +43,9 @@ "ListModelsResponse", "ListModelVersionsRequest", "ListModelVersionsResponse", + "ListModelVersionCheckpointsRequest", + "ModelVersionCheckpoint", + "ListModelVersionCheckpointsResponse", "UpdateModelRequest", "UpdateExplanationDatasetRequest", "UpdateExplanationDatasetOperationMetadata", @@ -391,6 +394,102 @@ def raw_page(self): ) +class ListModelVersionCheckpointsRequest(proto.Message): + r"""Request message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints]. + + Attributes: + name (str): + Required. The name of the model version to list checkpoints + for. + ``projects/{project}/locations/{location}/models/{model}@{version}`` + Example: + ``projects/{project}/locations/{location}/models/{model}@2`` + or + ``projects/{project}/locations/{location}/models/{model}@golden`` + If no version ID or alias is specified, the latest version + will be used. + page_size (int): + Optional. The standard list page size. + page_token (str): + Optional. The standard list page token. Typically obtained + via + [next_page_token][google.cloud.aiplatform.v1beta1.ListModelVersionCheckpointsResponse.next_page_token] + of the previous + [ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + call. + """ + + name: str = proto.Field( + proto.STRING, + number=1, + ) + page_size: int = proto.Field( + proto.INT32, + number=2, + ) + page_token: str = proto.Field( + proto.STRING, + number=3, + ) + + +class ModelVersionCheckpoint(proto.Message): + r"""A proto representation of a Spanner-stored + ModelVersionCheckpoint. The meaning of the fields is equivalent + to their in-Spanner counterparts. + + Attributes: + checkpoint_id (str): + The ID of the checkpoint. + epoch (int): + The epoch of the checkpoint. + step (int): + The step of the checkpoint. + """ + + checkpoint_id: str = proto.Field( + proto.STRING, + number=1, + ) + epoch: int = proto.Field( + proto.INT64, + number=2, + ) + step: int = proto.Field( + proto.INT64, + number=3, + ) + + +class ListModelVersionCheckpointsResponse(proto.Message): + r"""Response message for + [ModelService.ListModelVersionCheckpoints][google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints] + + Attributes: + checkpoints (MutableSequence[google.cloud.aiplatform_v1beta1.types.ModelVersionCheckpoint]): + List of Model Version checkpoints. + next_page_token (str): + A token to retrieve the next page of results. Pass to + [ListModelVersionCheckpointsRequest.page_token][google.cloud.aiplatform.v1beta1.ListModelVersionCheckpointsRequest.page_token] + to obtain that page. + """ + + @property + def raw_page(self): + return self + + checkpoints: MutableSequence["ModelVersionCheckpoint"] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message="ModelVersionCheckpoint", + ) + next_page_token: str = proto.Field( + proto.STRING, + number=2, + ) + + class UpdateModelRequest(proto.Message): r"""Request message for [ModelService.UpdateModel][google.cloud.aiplatform.v1beta1.ModelService.UpdateModel]. diff --git a/google/cloud/aiplatform_v1beta1/types/notebook_runtime.py b/google/cloud/aiplatform_v1beta1/types/notebook_runtime.py index 6b6b20b9d8..6b4895a520 100644 --- a/google/cloud/aiplatform_v1beta1/types/notebook_runtime.py +++ b/google/cloud/aiplatform_v1beta1/types/notebook_runtime.py @@ -27,6 +27,7 @@ from google.cloud.aiplatform_v1beta1.types import ( notebook_runtime_template_ref as gca_notebook_runtime_template_ref, ) +from google.cloud.aiplatform_v1beta1.types import notebook_software_config from google.protobuf import timestamp_pb2 # type: ignore @@ -150,6 +151,9 @@ class NotebookRuntimeTemplate(proto.Message): encryption_spec (google.cloud.aiplatform_v1beta1.types.EncryptionSpec): Customer-managed encryption key spec for the notebook runtime. + software_config (google.cloud.aiplatform_v1beta1.types.NotebookSoftwareConfig): + Optional. The notebook software configuration + of the notebook runtime. """ name: str = proto.Field( @@ -237,6 +241,11 @@ class NotebookRuntimeTemplate(proto.Message): number=23, message=gca_encryption_spec.EncryptionSpec, ) + software_config: notebook_software_config.NotebookSoftwareConfig = proto.Field( + proto.MESSAGE, + number=24, + message=notebook_software_config.NotebookSoftwareConfig, + ) class NotebookRuntime(proto.Message): @@ -351,6 +360,9 @@ class NotebookRuntime(proto.Message): Optional. The Compute Engine tags to add to runtime (see `Tagging instances `__). + software_config (google.cloud.aiplatform_v1beta1.types.NotebookSoftwareConfig): + Output only. Software config of the notebook + runtime. encryption_spec (google.cloud.aiplatform_v1beta1.types.EncryptionSpec): Output only. Customer-managed encryption key spec for the notebook runtime. @@ -521,6 +533,11 @@ class RuntimeState(proto.Enum): proto.STRING, number=25, ) + software_config: notebook_software_config.NotebookSoftwareConfig = proto.Field( + proto.MESSAGE, + number=31, + message=notebook_software_config.NotebookSoftwareConfig, + ) encryption_spec: gca_encryption_spec.EncryptionSpec = proto.Field( proto.MESSAGE, number=28, diff --git a/google/cloud/aiplatform_v1beta1/types/notebook_service.py b/google/cloud/aiplatform_v1beta1/types/notebook_service.py index c935b6875d..b86561ff0e 100644 --- a/google/cloud/aiplatform_v1beta1/types/notebook_service.py +++ b/google/cloud/aiplatform_v1beta1/types/notebook_service.py @@ -179,6 +179,8 @@ class ListNotebookRuntimeTemplatesRequest(proto.Message): - ``notebookRuntimeType`` supports = and !=. notebookRuntimeType enum: [USER_DEFINED, ONE_CLICK]. + - ``machineType`` supports = and !=. + - ``acceleratorType`` supports = and !=. Some examples: @@ -186,6 +188,8 @@ class ListNotebookRuntimeTemplatesRequest(proto.Message): - ``displayName="myDisplayName"`` - ``labels.myKey="myValue"`` - ``notebookRuntimeType=USER_DEFINED`` + - ``machineType=e2-standard-4`` + - ``acceleratorType=NVIDIA_TESLA_T4`` page_size (int): Optional. The standard list page size. page_token (str): @@ -441,6 +445,8 @@ class ListNotebookRuntimesRequest(proto.Message): UI_RESOURCE_STATE_CREATION_FAILED]. - ``notebookRuntimeType`` supports = and !=. notebookRuntimeType enum: [USER_DEFINED, ONE_CLICK]. + - ``machineType`` supports = and !=. + - ``acceleratorType`` supports = and !=. Some examples: @@ -453,6 +459,8 @@ class ListNotebookRuntimesRequest(proto.Message): - ``runtimeUser="test@google.com"`` - ``uiState=UI_RESOURCE_STATE_BEING_DELETED`` - ``notebookRuntimeType=USER_DEFINED`` + - ``machineType=e2-standard-4`` + - ``acceleratorType=NVIDIA_TESLA_T4`` page_size (int): Optional. The standard list page size. page_token (str): @@ -808,8 +816,9 @@ class ListNotebookExecutionJobsRequest(proto.Message): Optional. The standard list page size. page_token (str): Optional. The standard list page token. Typically obtained - via [ListNotebookExecutionJobs.next_page_token][] of the - previous + via + [ListNotebookExecutionJobsResponse.next_page_token][google.cloud.aiplatform.v1beta1.ListNotebookExecutionJobsResponse.next_page_token] + of the previous [NotebookService.ListNotebookExecutionJobs][google.cloud.aiplatform.v1beta1.NotebookService.ListNotebookExecutionJobs] call. order_by (str): @@ -863,8 +872,8 @@ class ListNotebookExecutionJobsResponse(proto.Message): requested page. next_page_token (str): A token to retrieve next page of results. Pass to - [ListNotebookExecutionJobs.page_token][] to obtain that - page. + [ListNotebookExecutionJobsRequest.page_token][google.cloud.aiplatform.v1beta1.ListNotebookExecutionJobsRequest.page_token] + to obtain that page. """ @property diff --git a/google/cloud/aiplatform_v1beta1/types/notebook_software_config.py b/google/cloud/aiplatform_v1beta1/types/notebook_software_config.py new file mode 100644 index 0000000000..4a6d530aef --- /dev/null +++ b/google/cloud/aiplatform_v1beta1/types/notebook_software_config.py @@ -0,0 +1,108 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from typing import MutableMapping, MutableSequence + +import proto # type: ignore + +from google.cloud.aiplatform_v1beta1.types import env_var + + +__protobuf__ = proto.module( + package="google.cloud.aiplatform.v1beta1", + manifest={ + "PostStartupScriptConfig", + "NotebookSoftwareConfig", + }, +) + + +class PostStartupScriptConfig(proto.Message): + r"""Post startup script config. + + Attributes: + post_startup_script (str): + Optional. Post startup script to run after + runtime is started. + post_startup_script_url (str): + Optional. Post startup script url to + download. Example: https://bucket/script.sh + post_startup_script_behavior (google.cloud.aiplatform_v1beta1.types.PostStartupScriptConfig.PostStartupScriptBehavior): + Optional. Post startup script behavior that + defines download and execution behavior. + """ + + class PostStartupScriptBehavior(proto.Enum): + r"""Represents a notebook runtime post startup script behavior. + + Values: + POST_STARTUP_SCRIPT_BEHAVIOR_UNSPECIFIED (0): + Unspecified post startup script behavior. + RUN_ONCE (1): + Run post startup script after runtime is + started. + RUN_EVERY_START (2): + Run post startup script after runtime is + stopped. + DOWNLOAD_AND_RUN_EVERY_START (3): + Download and run post startup script every + time runtime is started. + """ + POST_STARTUP_SCRIPT_BEHAVIOR_UNSPECIFIED = 0 + RUN_ONCE = 1 + RUN_EVERY_START = 2 + DOWNLOAD_AND_RUN_EVERY_START = 3 + + post_startup_script: str = proto.Field( + proto.STRING, + number=1, + ) + post_startup_script_url: str = proto.Field( + proto.STRING, + number=2, + ) + post_startup_script_behavior: PostStartupScriptBehavior = proto.Field( + proto.ENUM, + number=3, + enum=PostStartupScriptBehavior, + ) + + +class NotebookSoftwareConfig(proto.Message): + r"""Notebook Software Config. + + Attributes: + env (MutableSequence[google.cloud.aiplatform_v1beta1.types.EnvVar]): + Optional. Environment variables to be passed + to the container. Maximum limit is 100. + post_startup_script_config (google.cloud.aiplatform_v1beta1.types.PostStartupScriptConfig): + Optional. Post startup script config. + """ + + env: MutableSequence[env_var.EnvVar] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message=env_var.EnvVar, + ) + post_startup_script_config: "PostStartupScriptConfig" = proto.Field( + proto.MESSAGE, + number=2, + message="PostStartupScriptConfig", + ) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/google/cloud/aiplatform_v1beta1/types/prediction_service.py b/google/cloud/aiplatform_v1beta1/types/prediction_service.py index 64c275250a..6f8a1cf0c5 100644 --- a/google/cloud/aiplatform_v1beta1/types/prediction_service.py +++ b/google/cloud/aiplatform_v1beta1/types/prediction_service.py @@ -25,6 +25,7 @@ from google.cloud.aiplatform_v1beta1.types import tool from google.cloud.aiplatform_v1beta1.types import types from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore __protobuf__ = proto.module( @@ -961,6 +962,12 @@ class GenerateContentResponse(proto.Message): model_version (str): Output only. The model version used to generate the response. + create_time (google.protobuf.timestamp_pb2.Timestamp): + Output only. Timestamp when the request is + made to the server. + response_id (str): + Output only. response_id is used to identify each response. + It is the encoding of the event_id. prompt_feedback (google.cloud.aiplatform_v1beta1.types.GenerateContentResponse.PromptFeedback): Output only. Content filter results for a prompt sent in the request. Note: Sent only in @@ -1096,6 +1103,15 @@ class UsageMetadata(proto.Message): proto.STRING, number=11, ) + create_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=12, + message=timestamp_pb2.Timestamp, + ) + response_id: str = proto.Field( + proto.STRING, + number=13, + ) prompt_feedback: PromptFeedback = proto.Field( proto.MESSAGE, number=3, diff --git a/google/cloud/aiplatform_v1beta1/types/vertex_rag_data.py b/google/cloud/aiplatform_v1beta1/types/vertex_rag_data.py index da4a358158..b810c7514e 100644 --- a/google/cloud/aiplatform_v1beta1/types/vertex_rag_data.py +++ b/google/cloud/aiplatform_v1beta1/types/vertex_rag_data.py @@ -503,6 +503,9 @@ class RagCorpus(proto.Message): Vertex AI Search. This field is a member of `oneof`_ ``backend_config``. + rag_files_count (int): + Output only. The number of RagFiles in the + RagCorpus. """ name: str = proto.Field( @@ -554,6 +557,10 @@ class RagCorpus(proto.Message): oneof="backend_config", message="VertexAiSearchConfig", ) + rag_files_count: int = proto.Field( + proto.INT32, + number=11, + ) class RagFile(proto.Message): diff --git a/pypi/_vertex_ai_placeholder/version.py b/pypi/_vertex_ai_placeholder/version.py index d345120a24..5f357dd9b7 100644 --- a/pypi/_vertex_ai_placeholder/version.py +++ b/pypi/_vertex_ai_placeholder/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.80.0" +__version__ = "1.81.0" diff --git a/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_async.py b/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_async.py new file mode 100644 index 0000000000..a26da545d7 --- /dev/null +++ b/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_async.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for ListModelVersionCheckpoints +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_async] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1 + + +async def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1.ModelServiceAsyncClient() + + # Initialize request argument(s) + request = aiplatform_v1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + async for response in page_result: + print(response) + +# [END aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_async] diff --git a/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_sync.py b/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_sync.py new file mode 100644 index 0000000000..ca1bdf1ee4 --- /dev/null +++ b/samples/generated_samples/aiplatform_v1_generated_model_service_list_model_version_checkpoints_sync.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for ListModelVersionCheckpoints +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_sync] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1 + + +def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1.ModelServiceClient() + + # Initialize request argument(s) + request = aiplatform_v1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + for response in page_result: + print(response) + +# [END aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_sync] diff --git a/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_async.py b/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_async.py new file mode 100644 index 0000000000..be7b219d59 --- /dev/null +++ b/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_async.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for EvaluateDataset +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_async] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1beta1 + + +async def sample_evaluate_dataset(): + # Create a client + client = aiplatform_v1beta1.EvaluationServiceAsyncClient() + + # Initialize request argument(s) + dataset = aiplatform_v1beta1.EvaluationDataset() + dataset.gcs_source.uris = ['uris_value1', 'uris_value2'] + + output_config = aiplatform_v1beta1.OutputConfig() + output_config.gcs_destination.output_uri_prefix = "output_uri_prefix_value" + + request = aiplatform_v1beta1.EvaluateDatasetRequest( + location="location_value", + dataset=dataset, + output_config=output_config, + ) + + # Make the request + operation = client.evaluate_dataset(request=request) + + print("Waiting for operation to complete...") + + response = (await operation).result() + + # Handle the response + print(response) + +# [END aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_async] diff --git a/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_sync.py b/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_sync.py new file mode 100644 index 0000000000..ba73329837 --- /dev/null +++ b/samples/generated_samples/aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_sync.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for EvaluateDataset +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_sync] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1beta1 + + +def sample_evaluate_dataset(): + # Create a client + client = aiplatform_v1beta1.EvaluationServiceClient() + + # Initialize request argument(s) + dataset = aiplatform_v1beta1.EvaluationDataset() + dataset.gcs_source.uris = ['uris_value1', 'uris_value2'] + + output_config = aiplatform_v1beta1.OutputConfig() + output_config.gcs_destination.output_uri_prefix = "output_uri_prefix_value" + + request = aiplatform_v1beta1.EvaluateDatasetRequest( + location="location_value", + dataset=dataset, + output_config=output_config, + ) + + # Make the request + operation = client.evaluate_dataset(request=request) + + print("Waiting for operation to complete...") + + response = operation.result() + + # Handle the response + print(response) + +# [END aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_sync] diff --git a/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_async.py b/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_async.py new file mode 100644 index 0000000000..e690607d8e --- /dev/null +++ b/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_async.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for ListModelVersionCheckpoints +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_async] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1beta1 + + +async def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1beta1.ModelServiceAsyncClient() + + # Initialize request argument(s) + request = aiplatform_v1beta1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + async for response in page_result: + print(response) + +# [END aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_async] diff --git a/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_sync.py b/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_sync.py new file mode 100644 index 0000000000..98f0777f59 --- /dev/null +++ b/samples/generated_samples/aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_sync.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for ListModelVersionCheckpoints +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-aiplatform + + +# [START aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_sync] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import aiplatform_v1beta1 + + +def sample_list_model_version_checkpoints(): + # Create a client + client = aiplatform_v1beta1.ModelServiceClient() + + # Initialize request argument(s) + request = aiplatform_v1beta1.ListModelVersionCheckpointsRequest( + name="name_value", + ) + + # Make the request + page_result = client.list_model_version_checkpoints(request=request) + + # Handle the response + for response in page_result: + print(response) + +# [END aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_sync] diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json index f1c28dec17..6f95da4f4b 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.80.0" + "version": "1.81.0" }, "snippets": [ { @@ -32201,6 +32201,167 @@ ], "title": "aiplatform_v1_generated_model_service_list_model_evaluations_sync.py" }, + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.aiplatform_v1.ModelServiceAsyncClient", + "shortName": "ModelServiceAsyncClient" + }, + "fullName": "google.cloud.aiplatform_v1.ModelServiceAsyncClient.list_model_version_checkpoints", + "method": { + "fullName": "google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints", + "service": { + "fullName": "google.cloud.aiplatform.v1.ModelService", + "shortName": "ModelService" + }, + "shortName": "ListModelVersionCheckpoints" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.cloud.aiplatform_v1.services.model_service.pagers.ListModelVersionCheckpointsAsyncPager", + "shortName": "list_model_version_checkpoints" + }, + "description": "Sample for ListModelVersionCheckpoints", + "file": "aiplatform_v1_generated_model_service_list_model_version_checkpoints_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_async", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1_generated_model_service_list_model_version_checkpoints_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.aiplatform_v1.ModelServiceClient", + "shortName": "ModelServiceClient" + }, + "fullName": "google.cloud.aiplatform_v1.ModelServiceClient.list_model_version_checkpoints", + "method": { + "fullName": "google.cloud.aiplatform.v1.ModelService.ListModelVersionCheckpoints", + "service": { + "fullName": "google.cloud.aiplatform.v1.ModelService", + "shortName": "ModelService" + }, + "shortName": "ListModelVersionCheckpoints" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1.types.ListModelVersionCheckpointsRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.cloud.aiplatform_v1.services.model_service.pagers.ListModelVersionCheckpointsPager", + "shortName": "list_model_version_checkpoints" + }, + "description": "Sample for ListModelVersionCheckpoints", + "file": "aiplatform_v1_generated_model_service_list_model_version_checkpoints_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1_generated_ModelService_ListModelVersionCheckpoints_sync", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1_generated_model_service_list_model_version_checkpoints_sync.py" + }, { "canonical": true, "clientMethod": { diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json index c4fd68d40e..e2b2446a7f 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.80.0" + "version": "1.81.0" }, "snippets": [ { @@ -5621,6 +5621,159 @@ ], "title": "aiplatform_v1beta1_generated_endpoint_service_update_endpoint_sync.py" }, + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.aiplatform_v1beta1.EvaluationServiceAsyncClient", + "shortName": "EvaluationServiceAsyncClient" + }, + "fullName": "google.cloud.aiplatform_v1beta1.EvaluationServiceAsyncClient.evaluate_dataset", + "method": { + "fullName": "google.cloud.aiplatform.v1beta1.EvaluationService.EvaluateDataset", + "service": { + "fullName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "shortName": "EvaluationService" + }, + "shortName": "EvaluateDataset" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1beta1.types.EvaluateDatasetRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.api_core.operation_async.AsyncOperation", + "shortName": "evaluate_dataset" + }, + "description": "Sample for EvaluateDataset", + "file": "aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_async", + "segments": [ + { + "end": 63, + "start": 27, + "type": "FULL" + }, + { + "end": 63, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 53, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 60, + "start": 54, + "type": "REQUEST_EXECUTION" + }, + { + "end": 64, + "start": 61, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.aiplatform_v1beta1.EvaluationServiceClient", + "shortName": "EvaluationServiceClient" + }, + "fullName": "google.cloud.aiplatform_v1beta1.EvaluationServiceClient.evaluate_dataset", + "method": { + "fullName": "google.cloud.aiplatform.v1beta1.EvaluationService.EvaluateDataset", + "service": { + "fullName": "google.cloud.aiplatform.v1beta1.EvaluationService", + "shortName": "EvaluationService" + }, + "shortName": "EvaluateDataset" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1beta1.types.EvaluateDatasetRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.api_core.operation.Operation", + "shortName": "evaluate_dataset" + }, + "description": "Sample for EvaluateDataset", + "file": "aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1beta1_generated_EvaluationService_EvaluateDataset_sync", + "segments": [ + { + "end": 63, + "start": 27, + "type": "FULL" + }, + { + "end": 63, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 53, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 60, + "start": 54, + "type": "REQUEST_EXECUTION" + }, + { + "end": 64, + "start": 61, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1beta1_generated_evaluation_service_evaluate_dataset_sync.py" + }, { "canonical": true, "clientMethod": { @@ -36781,6 +36934,167 @@ ], "title": "aiplatform_v1beta1_generated_model_service_list_model_evaluations_sync.py" }, + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.aiplatform_v1beta1.ModelServiceAsyncClient", + "shortName": "ModelServiceAsyncClient" + }, + "fullName": "google.cloud.aiplatform_v1beta1.ModelServiceAsyncClient.list_model_version_checkpoints", + "method": { + "fullName": "google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints", + "service": { + "fullName": "google.cloud.aiplatform.v1beta1.ModelService", + "shortName": "ModelService" + }, + "shortName": "ListModelVersionCheckpoints" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.cloud.aiplatform_v1beta1.services.model_service.pagers.ListModelVersionCheckpointsAsyncPager", + "shortName": "list_model_version_checkpoints" + }, + "description": "Sample for ListModelVersionCheckpoints", + "file": "aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_async", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.aiplatform_v1beta1.ModelServiceClient", + "shortName": "ModelServiceClient" + }, + "fullName": "google.cloud.aiplatform_v1beta1.ModelServiceClient.list_model_version_checkpoints", + "method": { + "fullName": "google.cloud.aiplatform.v1beta1.ModelService.ListModelVersionCheckpoints", + "service": { + "fullName": "google.cloud.aiplatform.v1beta1.ModelService", + "shortName": "ModelService" + }, + "shortName": "ListModelVersionCheckpoints" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.aiplatform_v1beta1.types.ListModelVersionCheckpointsRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, Union[str, bytes]]]" + } + ], + "resultType": "google.cloud.aiplatform_v1beta1.services.model_service.pagers.ListModelVersionCheckpointsPager", + "shortName": "list_model_version_checkpoints" + }, + "description": "Sample for ListModelVersionCheckpoints", + "file": "aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "aiplatform_v1beta1_generated_ModelService_ListModelVersionCheckpoints_sync", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "aiplatform_v1beta1_generated_model_service_list_model_version_checkpoints_sync.py" + }, { "canonical": true, "clientMethod": { diff --git a/setup.py b/setup.py index 5738d4c3a5..848af1a267 100644 --- a/setup.py +++ b/setup.py @@ -156,9 +156,9 @@ ] langchain_extra_require = [ - "langchain >= 0.1.16, < 0.4", - "langchain-core < 0.4", - "langchain-google-vertexai < 3", + "langchain >= 0.3, < 0.4", + "langchain-core >= 0.3, < 0.4", + "langchain-google-vertexai >= 2, < 3", "langgraph >= 0.2.45, < 0.3", "openinference-instrumentation-langchain >= 0.1.19, < 0.2", ] diff --git a/tests/unit/aiplatform/test_vision_models.py b/tests/unit/aiplatform/test_vision_models.py index 20052a617d..64e060d974 100644 --- a/tests/unit/aiplatform/test_vision_models.py +++ b/tests/unit/aiplatform/test_vision_models.py @@ -1280,8 +1280,27 @@ def test_get_image_segmentation_results(self): target=prediction_service_client.PredictionServiceClient, attribute="predict", return_value=gca_prediction_response, - ): - segmentation_response = model.segment_image(base_image=image) + ) as mock_predict: + binary_color_threshold = 48 + segmentation_response = model.segment_image( + base_image=image, + confidence_threshold=0.1, + binary_color_threshold=binary_color_threshold, + ) + mock_predict.assert_called_once_with( + endpoint="projects/123456789/locations/us-central1/publishers/google/models/image-segmentation-001", + instances=[ + { + "base64EncodedImage": image._image_bytes, + "parameters": { + "mode": "foreground", + "confidenceThreshold": 0.1, + "binaryColorThreshold": 48, + }, + } + ], + retry=base._DEFAULT_RETRY, + ) assert len(segmentation_response) == 1 diff --git a/tests/unit/gapic/aiplatform_v1/test_dataset_service.py b/tests/unit/gapic/aiplatform_v1/test_dataset_service.py index 4c65039b21..9342e5a627 100644 --- a/tests/unit/gapic/aiplatform_v1/test_dataset_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_dataset_service.py @@ -97,6 +97,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -354,6 +362,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = DatasetServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = DatasetServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -14080,10 +14131,13 @@ def test_create_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_create_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_create_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_create_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetRequest.pb( dataset_service.CreateDatasetRequest() ) @@ -14107,6 +14161,7 @@ def test_create_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_dataset( request, @@ -14118,6 +14173,7 @@ def test_create_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_dataset_rest_bad_request(request_type=dataset_service.GetDatasetRequest): @@ -14218,10 +14274,13 @@ def test_get_dataset_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_get_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetRequest.pb( dataset_service.GetDatasetRequest() ) @@ -14245,6 +14304,7 @@ def test_get_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset.Dataset() + post_with_metadata.return_value = dataset.Dataset(), metadata client.get_dataset( request, @@ -14256,6 +14316,7 @@ def test_get_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_dataset_rest_bad_request( @@ -14467,10 +14528,13 @@ def test_update_dataset_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_update_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_update_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_update_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetRequest.pb( dataset_service.UpdateDatasetRequest() ) @@ -14494,6 +14558,7 @@ def test_update_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset.Dataset() + post_with_metadata.return_value = gca_dataset.Dataset(), metadata client.update_dataset( request, @@ -14505,6 +14570,7 @@ def test_update_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_datasets_rest_bad_request( @@ -14589,10 +14655,13 @@ def test_list_datasets_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_datasets" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_datasets_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_datasets" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetsRequest.pb( dataset_service.ListDatasetsRequest() ) @@ -14618,6 +14687,10 @@ def test_list_datasets_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetsResponse(), + metadata, + ) client.list_datasets( request, @@ -14629,6 +14702,7 @@ def test_list_datasets_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_dataset_rest_bad_request( @@ -14709,10 +14783,13 @@ def test_delete_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_delete_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetRequest.pb( dataset_service.DeleteDatasetRequest() ) @@ -14736,6 +14813,7 @@ def test_delete_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_dataset( request, @@ -14747,6 +14825,7 @@ def test_delete_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_data_rest_bad_request(request_type=dataset_service.ImportDataRequest): @@ -14825,10 +14904,13 @@ def test_import_data_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_import_data" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_import_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_import_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ImportDataRequest.pb( dataset_service.ImportDataRequest() ) @@ -14852,6 +14934,7 @@ def test_import_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_data( request, @@ -14863,6 +14946,7 @@ def test_import_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_data_rest_bad_request(request_type=dataset_service.ExportDataRequest): @@ -14941,10 +15025,13 @@ def test_export_data_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_export_data" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_export_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_export_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ExportDataRequest.pb( dataset_service.ExportDataRequest() ) @@ -14968,6 +15055,7 @@ def test_export_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_data( request, @@ -14979,6 +15067,7 @@ def test_export_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_dataset_version_rest_bad_request( @@ -15147,10 +15236,14 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_create_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_create_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_create_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetVersionRequest.pb( dataset_service.CreateDatasetVersionRequest() ) @@ -15174,6 +15267,7 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_dataset_version( request, @@ -15185,6 +15279,7 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_dataset_version_rest_bad_request( @@ -15377,10 +15472,14 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_update_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_update_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_update_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetVersionRequest.pb( dataset_service.UpdateDatasetVersionRequest() ) @@ -15406,6 +15505,7 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset_version.DatasetVersion() + post_with_metadata.return_value = gca_dataset_version.DatasetVersion(), metadata client.update_dataset_version( request, @@ -15417,6 +15517,7 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_dataset_version_rest_bad_request( @@ -15501,10 +15602,14 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_delete_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetVersionRequest.pb( dataset_service.DeleteDatasetVersionRequest() ) @@ -15528,6 +15633,7 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_dataset_version( request, @@ -15539,6 +15645,7 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_dataset_version_rest_bad_request( @@ -15639,10 +15746,14 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_get_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetVersionRequest.pb( dataset_service.GetDatasetVersionRequest() ) @@ -15668,6 +15779,7 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_version.DatasetVersion() + post_with_metadata.return_value = dataset_version.DatasetVersion(), metadata client.get_dataset_version( request, @@ -15679,6 +15791,7 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_dataset_versions_rest_bad_request( @@ -15763,10 +15876,14 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_dataset_versions" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_list_dataset_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_dataset_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetVersionsRequest.pb( dataset_service.ListDatasetVersionsRequest() ) @@ -15792,6 +15909,10 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetVersionsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetVersionsResponse(), + metadata, + ) client.list_dataset_versions( request, @@ -15803,6 +15924,7 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_restore_dataset_version_rest_bad_request( @@ -15887,10 +16009,14 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_restore_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_restore_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_restore_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.RestoreDatasetVersionRequest.pb( dataset_service.RestoreDatasetVersionRequest() ) @@ -15914,6 +16040,7 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.restore_dataset_version( request, @@ -15925,6 +16052,7 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_data_items_rest_bad_request( @@ -16009,10 +16137,13 @@ def test_list_data_items_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_data_items" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_data_items_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDataItemsRequest.pb( dataset_service.ListDataItemsRequest() ) @@ -16038,6 +16169,10 @@ def test_list_data_items_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDataItemsResponse(), + metadata, + ) client.list_data_items( request, @@ -16049,6 +16184,7 @@ def test_list_data_items_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_data_items_rest_bad_request( @@ -16133,10 +16269,13 @@ def test_search_data_items_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_search_data_items" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_search_data_items_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_search_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.SearchDataItemsRequest.pb( dataset_service.SearchDataItemsRequest() ) @@ -16162,6 +16301,10 @@ def test_search_data_items_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.SearchDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.SearchDataItemsResponse(), + metadata, + ) client.search_data_items( request, @@ -16173,6 +16316,7 @@ def test_search_data_items_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_saved_queries_rest_bad_request( @@ -16257,10 +16401,14 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_saved_queries" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_list_saved_queries_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_saved_queries" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListSavedQueriesRequest.pb( dataset_service.ListSavedQueriesRequest() ) @@ -16286,6 +16434,10 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListSavedQueriesResponse() + post_with_metadata.return_value = ( + dataset_service.ListSavedQueriesResponse(), + metadata, + ) client.list_saved_queries( request, @@ -16297,6 +16449,7 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_saved_query_rest_bad_request( @@ -16381,10 +16534,14 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_saved_query" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_delete_saved_query_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_saved_query" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteSavedQueryRequest.pb( dataset_service.DeleteSavedQueryRequest() ) @@ -16408,6 +16565,7 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_saved_query( request, @@ -16419,6 +16577,7 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_annotation_spec_rest_bad_request( @@ -16511,10 +16670,14 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_annotation_spec" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_get_annotation_spec_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_annotation_spec" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetAnnotationSpecRequest.pb( dataset_service.GetAnnotationSpecRequest() ) @@ -16540,6 +16703,7 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = annotation_spec.AnnotationSpec() + post_with_metadata.return_value = annotation_spec.AnnotationSpec(), metadata client.get_annotation_spec( request, @@ -16551,6 +16715,7 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_annotations_rest_bad_request( @@ -16639,10 +16804,13 @@ def test_list_annotations_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_annotations" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_annotations_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListAnnotationsRequest.pb( dataset_service.ListAnnotationsRequest() ) @@ -16668,6 +16836,10 @@ def test_list_annotations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListAnnotationsResponse() + post_with_metadata.return_value = ( + dataset_service.ListAnnotationsResponse(), + metadata, + ) client.list_annotations( request, @@ -16679,6 +16851,7 @@ def test_list_annotations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -17940,10 +18113,14 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_create_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_create_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_create_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetRequest.pb( dataset_service.CreateDatasetRequest() ) @@ -17967,6 +18144,7 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_dataset( request, @@ -17978,6 +18156,7 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18096,10 +18275,13 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetRequest.pb( dataset_service.GetDatasetRequest() ) @@ -18123,6 +18305,7 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset.Dataset() + post_with_metadata.return_value = dataset.Dataset(), metadata await client.get_dataset( request, @@ -18134,6 +18317,7 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18361,10 +18545,14 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_update_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_update_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_update_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetRequest.pb( dataset_service.UpdateDatasetRequest() ) @@ -18388,6 +18576,7 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset.Dataset() + post_with_metadata.return_value = gca_dataset.Dataset(), metadata await client.update_dataset( request, @@ -18399,6 +18588,7 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18499,10 +18689,14 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_datasets" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_datasets_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_datasets" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetsRequest.pb( dataset_service.ListDatasetsRequest() ) @@ -18528,6 +18722,10 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetsResponse(), + metadata, + ) await client.list_datasets( request, @@ -18539,6 +18737,7 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18635,10 +18834,14 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetRequest.pb( dataset_service.DeleteDatasetRequest() ) @@ -18662,6 +18865,7 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_dataset( request, @@ -18673,6 +18877,7 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18769,10 +18974,13 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_import_data" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_import_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_import_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ImportDataRequest.pb( dataset_service.ImportDataRequest() ) @@ -18796,6 +19004,7 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_data( request, @@ -18807,6 +19016,7 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18903,10 +19113,13 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_export_data" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_export_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_export_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ExportDataRequest.pb( dataset_service.ExportDataRequest() ) @@ -18930,6 +19143,7 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_data( request, @@ -18941,6 +19155,7 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19125,10 +19340,14 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_create_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_create_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_create_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetVersionRequest.pb( dataset_service.CreateDatasetVersionRequest() ) @@ -19152,6 +19371,7 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_dataset_version( request, @@ -19163,6 +19383,7 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19371,10 +19592,14 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_update_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_update_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_update_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetVersionRequest.pb( dataset_service.UpdateDatasetVersionRequest() ) @@ -19400,6 +19625,7 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_dataset_version.DatasetVersion() + post_with_metadata.return_value = gca_dataset_version.DatasetVersion(), metadata await client.update_dataset_version( request, @@ -19411,6 +19637,7 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19511,10 +19738,14 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetVersionRequest.pb( dataset_service.DeleteDatasetVersionRequest() ) @@ -19538,6 +19769,7 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_dataset_version( request, @@ -19549,6 +19781,7 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19665,10 +19898,14 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_get_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetVersionRequest.pb( dataset_service.GetDatasetVersionRequest() ) @@ -19694,6 +19931,7 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_version.DatasetVersion() + post_with_metadata.return_value = dataset_version.DatasetVersion(), metadata await client.get_dataset_version( request, @@ -19705,6 +19943,7 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19805,10 +20044,14 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_dataset_versions" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_dataset_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_dataset_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetVersionsRequest.pb( dataset_service.ListDatasetVersionsRequest() ) @@ -19834,6 +20077,10 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetVersionsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetVersionsResponse(), + metadata, + ) await client.list_dataset_versions( request, @@ -19845,6 +20092,7 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19945,10 +20193,14 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_restore_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_restore_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_restore_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.RestoreDatasetVersionRequest.pb( dataset_service.RestoreDatasetVersionRequest() ) @@ -19972,6 +20224,7 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.restore_dataset_version( request, @@ -19983,6 +20236,7 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20083,10 +20337,14 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_data_items" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_data_items_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDataItemsRequest.pb( dataset_service.ListDataItemsRequest() ) @@ -20112,6 +20370,10 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDataItemsResponse(), + metadata, + ) await client.list_data_items( request, @@ -20123,6 +20385,7 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20223,10 +20486,14 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_search_data_items" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_search_data_items_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_search_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.SearchDataItemsRequest.pb( dataset_service.SearchDataItemsRequest() ) @@ -20252,6 +20519,10 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.SearchDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.SearchDataItemsResponse(), + metadata, + ) await client.search_data_items( request, @@ -20263,6 +20534,7 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20363,10 +20635,14 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_saved_queries" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_saved_queries_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_saved_queries" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListSavedQueriesRequest.pb( dataset_service.ListSavedQueriesRequest() ) @@ -20392,6 +20668,10 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListSavedQueriesResponse() + post_with_metadata.return_value = ( + dataset_service.ListSavedQueriesResponse(), + metadata, + ) await client.list_saved_queries( request, @@ -20403,6 +20683,7 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20503,10 +20784,14 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_saved_query" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_saved_query_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_saved_query" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteSavedQueryRequest.pb( dataset_service.DeleteSavedQueryRequest() ) @@ -20530,6 +20815,7 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_saved_query( request, @@ -20541,6 +20827,7 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20649,10 +20936,14 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_annotation_spec" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_get_annotation_spec_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_annotation_spec" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetAnnotationSpecRequest.pb( dataset_service.GetAnnotationSpecRequest() ) @@ -20678,6 +20969,7 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = annotation_spec.AnnotationSpec() + post_with_metadata.return_value = annotation_spec.AnnotationSpec(), metadata await client.get_annotation_spec( request, @@ -20689,6 +20981,7 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20793,10 +21086,14 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_annotations" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListAnnotationsRequest.pb( dataset_service.ListAnnotationsRequest() ) @@ -20822,6 +21119,10 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListAnnotationsResponse() + post_with_metadata.return_value = ( + dataset_service.ListAnnotationsResponse(), + metadata, + ) await client.list_annotations( request, @@ -20833,6 +21134,7 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_deployment_resource_pool_service.py b/tests/unit/gapic/aiplatform_v1/test_deployment_resource_pool_service.py index 386e94fd76..6c5af6fcf1 100644 --- a/tests/unit/gapic/aiplatform_v1/test_deployment_resource_pool_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_deployment_resource_pool_service.py @@ -99,6 +99,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -385,6 +393,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = DeploymentResourcePoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = DeploymentResourcePoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5656,11 +5707,15 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_create_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_create_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_create_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest() @@ -5686,6 +5741,7 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_deployment_resource_pool( request, @@ -5697,6 +5753,7 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_deployment_resource_pool_rest_bad_request( @@ -5794,11 +5851,15 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_get_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_get_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_get_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.GetDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.GetDeploymentResourcePoolRequest() @@ -5826,6 +5887,10 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = deployment_resource_pool.DeploymentResourcePool() + post_with_metadata.return_value = ( + deployment_resource_pool.DeploymentResourcePool(), + metadata, + ) client.get_deployment_resource_pool( request, @@ -5837,6 +5902,7 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_deployment_resource_pools_rest_bad_request( @@ -5928,11 +5994,15 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_list_deployment_resource_pools", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_list_deployment_resource_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_list_deployment_resource_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest.pb( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest() @@ -5962,6 +6032,10 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): post.return_value = ( deployment_resource_pool_service.ListDeploymentResourcePoolsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse(), + metadata, + ) client.list_deployment_resource_pools( request, @@ -5973,6 +6047,7 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_deployment_resource_pool_rest_bad_request( @@ -6162,11 +6237,15 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_update_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_update_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_update_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest() @@ -6192,6 +6271,7 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_deployment_resource_pool( request, @@ -6203,6 +6283,7 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_deployment_resource_pool_rest_bad_request( @@ -6288,11 +6369,15 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_delete_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_delete_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_delete_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest() @@ -6318,6 +6403,7 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_deployment_resource_pool( request, @@ -6329,6 +6415,7 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_deployed_models_rest_bad_request( @@ -6424,11 +6511,15 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_query_deployed_models", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_query_deployed_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_query_deployed_models", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = deployment_resource_pool_service.QueryDeployedModelsRequest.pb( deployment_resource_pool_service.QueryDeployedModelsRequest() ) @@ -6458,6 +6549,10 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): post.return_value = ( deployment_resource_pool_service.QueryDeployedModelsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.QueryDeployedModelsResponse(), + metadata, + ) client.query_deployed_models( request, @@ -6469,6 +6564,7 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7370,11 +7466,15 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_create_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_create_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_create_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest() @@ -7400,6 +7500,7 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_deployment_resource_pool( request, @@ -7411,6 +7512,7 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7524,11 +7626,15 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_get_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_get_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_get_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.GetDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.GetDeploymentResourcePoolRequest() @@ -7556,6 +7662,10 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = deployment_resource_pool.DeploymentResourcePool() + post_with_metadata.return_value = ( + deployment_resource_pool.DeploymentResourcePool(), + metadata, + ) await client.get_deployment_resource_pool( request, @@ -7567,6 +7677,7 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7676,11 +7787,15 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_list_deployment_resource_pools", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_list_deployment_resource_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_list_deployment_resource_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest.pb( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest() @@ -7710,6 +7825,10 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( post.return_value = ( deployment_resource_pool_service.ListDeploymentResourcePoolsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse(), + metadata, + ) await client.list_deployment_resource_pools( request, @@ -7721,6 +7840,7 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7928,11 +8048,15 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_update_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_update_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_update_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest() @@ -7958,6 +8082,7 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_deployment_resource_pool( request, @@ -7969,6 +8094,7 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8072,11 +8198,15 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_delete_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_delete_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_delete_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest() @@ -8102,6 +8232,7 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_deployment_resource_pool( request, @@ -8113,6 +8244,7 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8224,11 +8356,15 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_query_deployed_models", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_query_deployed_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_query_deployed_models", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = deployment_resource_pool_service.QueryDeployedModelsRequest.pb( deployment_resource_pool_service.QueryDeployedModelsRequest() ) @@ -8258,6 +8394,10 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) post.return_value = ( deployment_resource_pool_service.QueryDeployedModelsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.QueryDeployedModelsResponse(), + metadata, + ) await client.query_deployed_models( request, @@ -8269,6 +8409,7 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_endpoint_service.py b/tests/unit/gapic/aiplatform_v1/test_endpoint_service.py index 0a605e5549..8777d09abd 100644 --- a/tests/unit/gapic/aiplatform_v1/test_endpoint_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_endpoint_service.py @@ -98,6 +98,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -355,6 +363,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = EndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = EndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -7182,10 +7233,13 @@ def test_create_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_create_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_create_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_create_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.CreateEndpointRequest.pb( endpoint_service.CreateEndpointRequest() ) @@ -7209,6 +7263,7 @@ def test_create_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_endpoint( request, @@ -7220,6 +7275,7 @@ def test_create_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_endpoint_rest_bad_request( @@ -7327,10 +7383,13 @@ def test_get_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_get_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_get_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_get_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.GetEndpointRequest.pb( endpoint_service.GetEndpointRequest() ) @@ -7354,6 +7413,7 @@ def test_get_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint.Endpoint() + post_with_metadata.return_value = endpoint.Endpoint(), metadata client.get_endpoint( request, @@ -7365,6 +7425,7 @@ def test_get_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_endpoints_rest_bad_request( @@ -7449,10 +7510,13 @@ def test_list_endpoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_list_endpoints" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_list_endpoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_list_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.ListEndpointsRequest.pb( endpoint_service.ListEndpointsRequest() ) @@ -7478,6 +7542,10 @@ def test_list_endpoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint_service.ListEndpointsResponse() + post_with_metadata.return_value = ( + endpoint_service.ListEndpointsResponse(), + metadata, + ) client.list_endpoints( request, @@ -7489,6 +7557,7 @@ def test_list_endpoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_endpoint_rest_bad_request( @@ -7804,10 +7873,13 @@ def test_update_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_update_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_update_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_update_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointRequest.pb( endpoint_service.UpdateEndpointRequest() ) @@ -7831,6 +7903,7 @@ def test_update_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_endpoint.Endpoint() + post_with_metadata.return_value = gca_endpoint.Endpoint(), metadata client.update_endpoint( request, @@ -7842,6 +7915,7 @@ def test_update_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_endpoint_long_running_rest_bad_request( @@ -7926,10 +8000,14 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_update_endpoint_long_running" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, + "post_update_endpoint_long_running_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_update_endpoint_long_running" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointLongRunningRequest.pb( endpoint_service.UpdateEndpointLongRunningRequest() ) @@ -7953,6 +8031,7 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_endpoint_long_running( request, @@ -7964,6 +8043,7 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_endpoint_rest_bad_request( @@ -8044,10 +8124,13 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_delete_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_delete_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_delete_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeleteEndpointRequest.pb( endpoint_service.DeleteEndpointRequest() ) @@ -8071,6 +8154,7 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_endpoint( request, @@ -8082,6 +8166,7 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_deploy_model_rest_bad_request( @@ -8162,10 +8247,13 @@ def test_deploy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_deploy_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_deploy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_deploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeployModelRequest.pb( endpoint_service.DeployModelRequest() ) @@ -8189,6 +8277,7 @@ def test_deploy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.deploy_model( request, @@ -8200,6 +8289,7 @@ def test_deploy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_undeploy_model_rest_bad_request( @@ -8280,10 +8370,13 @@ def test_undeploy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_undeploy_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_undeploy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_undeploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UndeployModelRequest.pb( endpoint_service.UndeployModelRequest() ) @@ -8307,6 +8400,7 @@ def test_undeploy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.undeploy_model( request, @@ -8318,6 +8412,7 @@ def test_undeploy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_mutate_deployed_model_rest_bad_request( @@ -8398,10 +8493,14 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_mutate_deployed_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, + "post_mutate_deployed_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_mutate_deployed_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.MutateDeployedModelRequest.pb( endpoint_service.MutateDeployedModelRequest() ) @@ -8425,6 +8524,7 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.mutate_deployed_model( request, @@ -8436,6 +8536,7 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -9580,10 +9681,14 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_create_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_create_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_create_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.CreateEndpointRequest.pb( endpoint_service.CreateEndpointRequest() ) @@ -9607,6 +9712,7 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_endpoint( request, @@ -9618,6 +9724,7 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9741,10 +9848,14 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_get_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_get_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_get_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.GetEndpointRequest.pb( endpoint_service.GetEndpointRequest() ) @@ -9768,6 +9879,7 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint.Endpoint() + post_with_metadata.return_value = endpoint.Endpoint(), metadata await client.get_endpoint( request, @@ -9779,6 +9891,7 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9879,10 +9992,14 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_list_endpoints" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_list_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_list_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.ListEndpointsRequest.pb( endpoint_service.ListEndpointsRequest() ) @@ -9908,6 +10025,10 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint_service.ListEndpointsResponse() + post_with_metadata.return_value = ( + endpoint_service.ListEndpointsResponse(), + metadata, + ) await client.list_endpoints( request, @@ -9919,6 +10040,7 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10250,10 +10372,14 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_update_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_update_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_update_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointRequest.pb( endpoint_service.UpdateEndpointRequest() ) @@ -10277,6 +10403,7 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_endpoint.Endpoint() + post_with_metadata.return_value = gca_endpoint.Endpoint(), metadata await client.update_endpoint( request, @@ -10288,6 +10415,7 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10389,11 +10517,15 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter transports.AsyncEndpointServiceRestInterceptor, "post_update_endpoint_long_running", ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_update_endpoint_long_running_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_update_endpoint_long_running", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointLongRunningRequest.pb( endpoint_service.UpdateEndpointLongRunningRequest() ) @@ -10417,6 +10549,7 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_endpoint_long_running( request, @@ -10428,6 +10561,7 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10524,10 +10658,14 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_delete_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_delete_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_delete_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeleteEndpointRequest.pb( endpoint_service.DeleteEndpointRequest() ) @@ -10551,6 +10689,7 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_endpoint( request, @@ -10562,6 +10701,7 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10658,10 +10798,14 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_deploy_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_deploy_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_deploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeployModelRequest.pb( endpoint_service.DeployModelRequest() ) @@ -10685,6 +10829,7 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.deploy_model( request, @@ -10696,6 +10841,7 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10792,10 +10938,14 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_undeploy_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_undeploy_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_undeploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UndeployModelRequest.pb( endpoint_service.UndeployModelRequest() ) @@ -10819,6 +10969,7 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.undeploy_model( request, @@ -10830,6 +10981,7 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10926,10 +11078,14 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_mutate_deployed_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_mutate_deployed_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_mutate_deployed_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.MutateDeployedModelRequest.pb( endpoint_service.MutateDeployedModelRequest() ) @@ -10953,6 +11109,7 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.mutate_deployed_model( request, @@ -10964,6 +11121,7 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_evaluation_service.py b/tests/unit/gapic/aiplatform_v1/test_evaluation_service.py index a82fd0ce28..479bc99f71 100644 --- a/tests/unit/gapic/aiplatform_v1/test_evaluation_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_evaluation_service.py @@ -79,6 +79,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -345,6 +353,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = EvaluationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = EvaluationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1747,10 +1798,14 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EvaluationServiceRestInterceptor, "post_evaluate_instances" ) as post, mock.patch.object( + transports.EvaluationServiceRestInterceptor, + "post_evaluate_instances_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.EvaluationServiceRestInterceptor, "pre_evaluate_instances" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = evaluation_service.EvaluateInstancesRequest.pb( evaluation_service.EvaluateInstancesRequest() ) @@ -1776,6 +1831,10 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = evaluation_service.EvaluateInstancesResponse() + post_with_metadata.return_value = ( + evaluation_service.EvaluateInstancesResponse(), + metadata, + ) client.evaluate_instances( request, @@ -1787,6 +1846,7 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -2549,10 +2609,14 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEvaluationServiceRestInterceptor, "post_evaluate_instances" ) as post, mock.patch.object( + transports.AsyncEvaluationServiceRestInterceptor, + "post_evaluate_instances_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEvaluationServiceRestInterceptor, "pre_evaluate_instances" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = evaluation_service.EvaluateInstancesRequest.pb( evaluation_service.EvaluateInstancesRequest() ) @@ -2578,6 +2642,10 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = evaluation_service.EvaluateInstancesResponse() + post_with_metadata.return_value = ( + evaluation_service.EvaluateInstancesResponse(), + metadata, + ) await client.evaluate_instances( request, @@ -2589,6 +2657,7 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_feature_online_store_admin_service.py b/tests/unit/gapic/aiplatform_v1/test_feature_online_store_admin_service.py index 4fcf93a529..a9b9fbdb4f 100644 --- a/tests/unit/gapic/aiplatform_v1/test_feature_online_store_admin_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_feature_online_store_admin_service.py @@ -87,6 +87,7 @@ from google.cloud.aiplatform_v1.types import feature_view from google.cloud.aiplatform_v1.types import feature_view as gca_feature_view from google.cloud.aiplatform_v1.types import feature_view_sync +from google.cloud.aiplatform_v1.types import machine_resources from google.cloud.aiplatform_v1.types import operation as gca_operation from google.cloud.aiplatform_v1.types import service_networking from google.cloud.location import locations_pb2 @@ -103,6 +104,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -392,6 +401,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureOnlineStoreAdminServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureOnlineStoreAdminServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -10350,11 +10402,15 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest() @@ -10380,6 +10436,7 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_online_store( request, @@ -10391,6 +10448,7 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_online_store_rest_bad_request( @@ -10488,11 +10546,15 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.GetFeatureOnlineStoreRequest() ) @@ -10518,6 +10580,10 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store.FeatureOnlineStore() + post_with_metadata.return_value = ( + feature_online_store.FeatureOnlineStore(), + metadata, + ) client.get_feature_online_store( request, @@ -10529,6 +10595,7 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_online_stores_rest_bad_request( @@ -10620,11 +10687,15 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_online_stores", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_online_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_online_stores", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.ListFeatureOnlineStoresRequest.pb( feature_online_store_admin_service.ListFeatureOnlineStoresRequest() @@ -10656,6 +10727,10 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureOnlineStoresResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureOnlineStoresResponse(), + metadata, + ) client.list_feature_online_stores( request, @@ -10667,6 +10742,7 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_online_store_rest_bad_request( @@ -10860,11 +10936,15 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest() @@ -10890,6 +10970,7 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_online_store( request, @@ -10901,6 +10982,7 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_online_store_rest_bad_request( @@ -10986,11 +11068,15 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest() @@ -11016,6 +11102,7 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_online_store( request, @@ -11027,6 +11114,7 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_view_rest_bad_request( @@ -11105,6 +11193,12 @@ def test_create_feature_view_rest_call_success(request_type): "embedding_dimension": 1988, "distance_measure_type": 1, }, + "optimized_config": { + "automatic_resources": { + "min_replica_count": 1803, + "max_replica_count": 1805, + } + }, "service_agent_type": 1, "service_account_email": "service_account_email_value", "satisfies_pzs": True, @@ -11221,11 +11315,15 @@ def test_create_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.CreateFeatureViewRequest.pb( feature_online_store_admin_service.CreateFeatureViewRequest() ) @@ -11249,6 +11347,7 @@ def test_create_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_view( request, @@ -11260,6 +11359,7 @@ def test_create_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_view_rest_bad_request( @@ -11362,10 +11462,14 @@ def test_get_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewRequest.pb( feature_online_store_admin_service.GetFeatureViewRequest() ) @@ -11389,6 +11493,7 @@ def test_get_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view.FeatureView() + post_with_metadata.return_value = feature_view.FeatureView(), metadata client.get_feature_view( request, @@ -11400,6 +11505,7 @@ def test_get_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_views_rest_bad_request( @@ -11491,11 +11597,15 @@ def test_list_feature_views_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_views", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_views_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_views", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewsRequest.pb( feature_online_store_admin_service.ListFeatureViewsRequest() ) @@ -11525,6 +11635,10 @@ def test_list_feature_views_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewsResponse(), + metadata, + ) client.list_feature_views( request, @@ -11536,6 +11650,7 @@ def test_list_feature_views_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_view_rest_bad_request( @@ -11618,6 +11733,12 @@ def test_update_feature_view_rest_call_success(request_type): "embedding_dimension": 1988, "distance_measure_type": 1, }, + "optimized_config": { + "automatic_resources": { + "min_replica_count": 1803, + "max_replica_count": 1805, + } + }, "service_agent_type": 1, "service_account_email": "service_account_email_value", "satisfies_pzs": True, @@ -11734,11 +11855,15 @@ def test_update_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.UpdateFeatureViewRequest.pb( feature_online_store_admin_service.UpdateFeatureViewRequest() ) @@ -11762,6 +11887,7 @@ def test_update_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_view( request, @@ -11773,6 +11899,7 @@ def test_update_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_view_rest_bad_request( @@ -11858,11 +11985,15 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.DeleteFeatureViewRequest.pb( feature_online_store_admin_service.DeleteFeatureViewRequest() ) @@ -11886,6 +12017,7 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_view( request, @@ -11897,6 +12029,7 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_sync_feature_view_rest_bad_request( @@ -11990,11 +12123,15 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_sync_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_sync_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_sync_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.SyncFeatureViewRequest.pb( feature_online_store_admin_service.SyncFeatureViewRequest() ) @@ -12022,6 +12159,10 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_admin_service.SyncFeatureViewResponse() + post_with_metadata.return_value = ( + feature_online_store_admin_service.SyncFeatureViewResponse(), + metadata, + ) client.sync_feature_view( request, @@ -12033,6 +12174,7 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_view_sync_rest_bad_request( @@ -12126,11 +12268,15 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view_sync", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_sync_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view_sync", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewSyncRequest.pb( feature_online_store_admin_service.GetFeatureViewSyncRequest() ) @@ -12156,6 +12302,7 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view_sync.FeatureViewSync() + post_with_metadata.return_value = feature_view_sync.FeatureViewSync(), metadata client.get_feature_view_sync( request, @@ -12167,6 +12314,7 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_view_syncs_rest_bad_request( @@ -12260,11 +12408,15 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_view_syncs", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_view_syncs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_view_syncs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewSyncsRequest.pb( feature_online_store_admin_service.ListFeatureViewSyncsRequest() ) @@ -12294,6 +12446,10 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewSyncsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewSyncsResponse(), + metadata, + ) client.list_feature_view_syncs( request, @@ -12305,6 +12461,7 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -13458,11 +13615,15 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest() @@ -13488,6 +13649,7 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_online_store( request, @@ -13499,6 +13661,7 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13612,11 +13775,15 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.GetFeatureOnlineStoreRequest() ) @@ -13642,6 +13809,10 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = feature_online_store.FeatureOnlineStore() + post_with_metadata.return_value = ( + feature_online_store.FeatureOnlineStore(), + metadata, + ) await client.get_feature_online_store( request, @@ -13653,6 +13824,7 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13760,11 +13932,15 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_online_stores", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_online_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_online_stores", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.ListFeatureOnlineStoresRequest.pb( feature_online_store_admin_service.ListFeatureOnlineStoresRequest() @@ -13796,6 +13972,10 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce post.return_value = ( feature_online_store_admin_service.ListFeatureOnlineStoresResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureOnlineStoresResponse(), + metadata, + ) await client.list_feature_online_stores( request, @@ -13807,6 +13987,7 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14016,11 +14197,15 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest() @@ -14046,6 +14231,7 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_online_store( request, @@ -14057,6 +14243,7 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14158,11 +14345,15 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest() @@ -14188,6 +14379,7 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_online_store( request, @@ -14199,6 +14391,7 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14286,6 +14479,12 @@ async def test_create_feature_view_rest_asyncio_call_success(request_type): "embedding_dimension": 1988, "distance_measure_type": 1, }, + "optimized_config": { + "automatic_resources": { + "min_replica_count": 1803, + "max_replica_count": 1805, + } + }, "service_agent_type": 1, "service_account_email": "service_account_email_value", "satisfies_pzs": True, @@ -14409,11 +14608,15 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.CreateFeatureViewRequest.pb( feature_online_store_admin_service.CreateFeatureViewRequest() ) @@ -14437,6 +14640,7 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_view( request, @@ -14448,6 +14652,7 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14566,11 +14771,15 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewRequest.pb( feature_online_store_admin_service.GetFeatureViewRequest() ) @@ -14594,6 +14803,7 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view.FeatureView() + post_with_metadata.return_value = feature_view.FeatureView(), metadata await client.get_feature_view( request, @@ -14605,6 +14815,7 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14712,11 +14923,15 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_views", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_views_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_views", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewsRequest.pb( feature_online_store_admin_service.ListFeatureViewsRequest() ) @@ -14746,6 +14961,10 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewsResponse(), + metadata, + ) await client.list_feature_views( request, @@ -14757,6 +14976,7 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14848,6 +15068,12 @@ async def test_update_feature_view_rest_asyncio_call_success(request_type): "embedding_dimension": 1988, "distance_measure_type": 1, }, + "optimized_config": { + "automatic_resources": { + "min_replica_count": 1803, + "max_replica_count": 1805, + } + }, "service_agent_type": 1, "service_account_email": "service_account_email_value", "satisfies_pzs": True, @@ -14971,11 +15197,15 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.UpdateFeatureViewRequest.pb( feature_online_store_admin_service.UpdateFeatureViewRequest() ) @@ -14999,6 +15229,7 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_view( request, @@ -15010,6 +15241,7 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15111,11 +15343,15 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.DeleteFeatureViewRequest.pb( feature_online_store_admin_service.DeleteFeatureViewRequest() ) @@ -15139,6 +15375,7 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_view( request, @@ -15150,6 +15387,7 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15259,11 +15497,15 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_sync_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_sync_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_sync_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.SyncFeatureViewRequest.pb( feature_online_store_admin_service.SyncFeatureViewRequest() ) @@ -15291,6 +15533,10 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_admin_service.SyncFeatureViewResponse() + post_with_metadata.return_value = ( + feature_online_store_admin_service.SyncFeatureViewResponse(), + metadata, + ) await client.sync_feature_view( request, @@ -15302,6 +15548,7 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15411,11 +15658,15 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view_sync", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_sync_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view_sync", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewSyncRequest.pb( feature_online_store_admin_service.GetFeatureViewSyncRequest() ) @@ -15441,6 +15692,7 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = feature_view_sync.FeatureViewSync() + post_with_metadata.return_value = feature_view_sync.FeatureViewSync(), metadata await client.get_feature_view_sync( request, @@ -15452,6 +15704,7 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15561,11 +15814,15 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_view_syncs", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_view_syncs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_view_syncs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewSyncsRequest.pb( feature_online_store_admin_service.ListFeatureViewSyncsRequest() ) @@ -15595,6 +15852,10 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto post.return_value = ( feature_online_store_admin_service.ListFeatureViewSyncsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewSyncsResponse(), + metadata, + ) await client.list_feature_view_syncs( request, @@ -15606,6 +15867,7 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_feature_online_store_service.py b/tests/unit/gapic/aiplatform_v1/test_feature_online_store_service.py index 34324d87ab..aabb4da07d 100644 --- a/tests/unit/gapic/aiplatform_v1/test_feature_online_store_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_feature_online_store_service.py @@ -80,6 +80,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -358,6 +366,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureOnlineStoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureOnlineStoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2397,10 +2448,14 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "post_fetch_feature_values" ) as post, mock.patch.object( + transports.FeatureOnlineStoreServiceRestInterceptor, + "post_fetch_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "pre_fetch_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.FetchFeatureValuesRequest.pb( feature_online_store_service.FetchFeatureValuesRequest() ) @@ -2426,6 +2481,10 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.FetchFeatureValuesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.FetchFeatureValuesResponse(), + metadata, + ) client.fetch_feature_values( request, @@ -2437,6 +2496,7 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_nearest_entities_rest_bad_request( @@ -2527,11 +2587,15 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreServiceRestInterceptor, "post_search_nearest_entities", ) as post, mock.patch.object( + transports.FeatureOnlineStoreServiceRestInterceptor, + "post_search_nearest_entities_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "pre_search_nearest_entities", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.SearchNearestEntitiesRequest.pb( feature_online_store_service.SearchNearestEntitiesRequest() ) @@ -2559,6 +2623,10 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.SearchNearestEntitiesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.SearchNearestEntitiesResponse(), + metadata, + ) client.search_nearest_entities( request, @@ -2570,6 +2638,7 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3361,11 +3430,15 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "post_fetch_feature_values", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreServiceRestInterceptor, + "post_fetch_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "pre_fetch_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.FetchFeatureValuesRequest.pb( feature_online_store_service.FetchFeatureValuesRequest() ) @@ -3391,6 +3464,10 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.FetchFeatureValuesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.FetchFeatureValuesResponse(), + metadata, + ) await client.fetch_feature_values( request, @@ -3402,6 +3479,7 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3508,11 +3586,15 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "post_search_nearest_entities", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreServiceRestInterceptor, + "post_search_nearest_entities_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "pre_search_nearest_entities", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.SearchNearestEntitiesRequest.pb( feature_online_store_service.SearchNearestEntitiesRequest() ) @@ -3540,6 +3622,10 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = feature_online_store_service.SearchNearestEntitiesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.SearchNearestEntitiesResponse(), + metadata, + ) await client.search_nearest_entities( request, @@ -3551,6 +3637,7 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_feature_registry_service.py b/tests/unit/gapic/aiplatform_v1/test_feature_registry_service.py index 8d09591d55..ceb7ba2ebe 100644 --- a/tests/unit/gapic/aiplatform_v1/test_feature_registry_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_feature_registry_service.py @@ -95,6 +95,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -369,6 +377,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -8514,10 +8565,14 @@ def test_create_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_create_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureGroupRequest.pb( feature_registry_service.CreateFeatureGroupRequest() ) @@ -8541,6 +8596,7 @@ def test_create_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_group( request, @@ -8552,6 +8608,7 @@ def test_create_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_group_rest_bad_request( @@ -8640,10 +8697,14 @@ def test_get_feature_group_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureGroupRequest.pb( feature_registry_service.GetFeatureGroupRequest() ) @@ -8667,6 +8728,7 @@ def test_get_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_group.FeatureGroup() + post_with_metadata.return_value = feature_group.FeatureGroup(), metadata client.get_feature_group( request, @@ -8678,6 +8740,7 @@ def test_get_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_groups_rest_bad_request( @@ -8764,10 +8827,14 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_list_feature_groups" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_feature_groups_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_feature_groups" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureGroupsRequest.pb( feature_registry_service.ListFeatureGroupsRequest() ) @@ -8793,6 +8860,10 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureGroupsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureGroupsResponse(), + metadata, + ) client.list_feature_groups( request, @@ -8804,6 +8875,7 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_group_rest_bad_request( @@ -8979,10 +9051,14 @@ def test_update_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_update_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_update_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_update_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureGroupRequest.pb( feature_registry_service.UpdateFeatureGroupRequest() ) @@ -9006,6 +9082,7 @@ def test_update_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_group( request, @@ -9017,6 +9094,7 @@ def test_update_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_group_rest_bad_request( @@ -9097,10 +9175,14 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_delete_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_delete_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_delete_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureGroupRequest.pb( feature_registry_service.DeleteFeatureGroupRequest() ) @@ -9124,6 +9206,7 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_group( request, @@ -9135,6 +9218,7 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_rest_bad_request( @@ -9312,10 +9396,14 @@ def test_create_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -9339,6 +9427,7 @@ def test_create_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature( request, @@ -9350,6 +9439,7 @@ def test_create_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_features_rest_bad_request( @@ -9434,10 +9524,14 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -9461,6 +9555,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_create_features( request, @@ -9472,6 +9567,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_rest_bad_request( @@ -9572,10 +9668,14 @@ def test_get_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -9599,6 +9699,7 @@ def test_get_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata client.get_feature( request, @@ -9610,6 +9711,7 @@ def test_get_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_features_rest_bad_request( @@ -9698,10 +9800,14 @@ def test_list_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -9727,6 +9833,10 @@ def test_list_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) client.list_features( request, @@ -9738,6 +9848,7 @@ def test_list_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_rest_bad_request( @@ -9919,10 +10030,14 @@ def test_update_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -9946,6 +10061,7 @@ def test_update_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature( request, @@ -9957,6 +10073,7 @@ def test_update_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_rest_bad_request( @@ -10041,10 +10158,14 @@ def test_delete_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -10068,6 +10189,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature( request, @@ -10079,6 +10201,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -11155,11 +11278,15 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureGroupRequest.pb( feature_registry_service.CreateFeatureGroupRequest() ) @@ -11183,6 +11310,7 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_group( request, @@ -11194,6 +11322,7 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11298,10 +11427,14 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature_group" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureGroupRequest.pb( feature_registry_service.GetFeatureGroupRequest() ) @@ -11325,6 +11458,7 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_group.FeatureGroup() + post_with_metadata.return_value = feature_group.FeatureGroup(), metadata await client.get_feature_group( request, @@ -11336,6 +11470,7 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11439,10 +11574,14 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_feature_groups", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_feature_groups_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_feature_groups" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureGroupsRequest.pb( feature_registry_service.ListFeatureGroupsRequest() ) @@ -11468,6 +11607,10 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureGroupsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureGroupsResponse(), + metadata, + ) await client.list_feature_groups( request, @@ -11479,6 +11622,7 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11671,11 +11815,15 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_update_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_update_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_update_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureGroupRequest.pb( feature_registry_service.UpdateFeatureGroupRequest() ) @@ -11699,6 +11847,7 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_group( request, @@ -11710,6 +11859,7 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11807,11 +11957,15 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_delete_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_delete_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_delete_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureGroupRequest.pb( feature_registry_service.DeleteFeatureGroupRequest() ) @@ -11835,6 +11989,7 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_group( request, @@ -11846,6 +12001,7 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12039,10 +12195,14 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -12066,6 +12226,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature( request, @@ -12077,6 +12238,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12178,11 +12340,15 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) transports.AsyncFeatureRegistryServiceRestInterceptor, "post_batch_create_features", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_batch_create_features", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -12206,6 +12372,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_create_features( request, @@ -12217,6 +12384,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12333,10 +12501,14 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -12360,6 +12532,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata await client.get_feature( request, @@ -12371,6 +12544,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12475,10 +12649,14 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -12504,6 +12682,10 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) await client.list_features( request, @@ -12515,6 +12697,7 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12712,10 +12895,14 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -12739,6 +12926,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature( request, @@ -12750,6 +12938,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12850,10 +13039,14 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -12877,6 +13070,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature( request, @@ -12888,6 +13082,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_featurestore_online_serving_service.py b/tests/unit/gapic/aiplatform_v1/test_featurestore_online_serving_service.py index 678fde8481..830fbf6bed 100644 --- a/tests/unit/gapic/aiplatform_v1/test_featurestore_online_serving_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_featurestore_online_serving_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -377,6 +385,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeaturestoreOnlineServingServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeaturestoreOnlineServingServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3228,11 +3279,15 @@ def test_read_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_read_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.ReadFeatureValuesRequest.pb( featurestore_online_service.ReadFeatureValuesRequest() ) @@ -3258,6 +3313,10 @@ def test_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) client.read_feature_values( request, @@ -3269,6 +3328,7 @@ def test_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_streaming_read_feature_values_rest_bad_request( @@ -3361,11 +3421,15 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_streaming_read_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_streaming_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_streaming_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.StreamingReadFeatureValuesRequest.pb( featurestore_online_service.StreamingReadFeatureValuesRequest() ) @@ -3391,6 +3455,10 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) client.streaming_read_feature_values( request, @@ -3402,6 +3470,7 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_feature_values_rest_bad_request( @@ -3490,11 +3559,15 @@ def test_write_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_write_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_write_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_write_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.WriteFeatureValuesRequest.pb( featurestore_online_service.WriteFeatureValuesRequest() ) @@ -3520,6 +3593,10 @@ def test_write_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.WriteFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.WriteFeatureValuesResponse(), + metadata, + ) client.write_feature_values( request, @@ -3531,6 +3608,7 @@ def test_write_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -4344,11 +4422,15 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.ReadFeatureValuesRequest.pb( featurestore_online_service.ReadFeatureValuesRequest() ) @@ -4374,6 +4456,10 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) await client.read_feature_values( request, @@ -4385,6 +4471,7 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4493,11 +4580,15 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_streaming_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_streaming_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_streaming_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.StreamingReadFeatureValuesRequest.pb( featurestore_online_service.StreamingReadFeatureValuesRequest() ) @@ -4523,6 +4614,10 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) await client.streaming_read_feature_values( request, @@ -4534,6 +4629,7 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4638,11 +4734,15 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_write_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_write_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_write_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.WriteFeatureValuesRequest.pb( featurestore_online_service.WriteFeatureValuesRequest() ) @@ -4668,6 +4768,10 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.WriteFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.WriteFeatureValuesResponse(), + metadata, + ) await client.write_feature_values( request, @@ -4679,6 +4783,7 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_featurestore_service.py b/tests/unit/gapic/aiplatform_v1/test_featurestore_service.py index c805bfdcf3..9d9ffe7b95 100644 --- a/tests/unit/gapic/aiplatform_v1/test_featurestore_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_featurestore_service.py @@ -100,6 +100,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -372,6 +380,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeaturestoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeaturestoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -14950,10 +15001,14 @@ def test_create_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeaturestoreRequest.pb( featurestore_service.CreateFeaturestoreRequest() ) @@ -14977,6 +15032,7 @@ def test_create_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_featurestore( request, @@ -14988,6 +15044,7 @@ def test_create_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_featurestore_rest_bad_request( @@ -15082,10 +15139,14 @@ def test_get_featurestore_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_get_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeaturestoreRequest.pb( featurestore_service.GetFeaturestoreRequest() ) @@ -15109,6 +15170,7 @@ def test_get_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore.Featurestore() + post_with_metadata.return_value = featurestore.Featurestore(), metadata client.get_featurestore( request, @@ -15120,6 +15182,7 @@ def test_get_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_featurestores_rest_bad_request( @@ -15204,10 +15267,14 @@ def test_list_featurestores_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_featurestores" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_featurestores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_featurestores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturestoresRequest.pb( featurestore_service.ListFeaturestoresRequest() ) @@ -15233,6 +15300,10 @@ def test_list_featurestores_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturestoresResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturestoresResponse(), + metadata, + ) client.list_featurestores( request, @@ -15244,6 +15315,7 @@ def test_list_featurestores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_featurestore_rest_bad_request( @@ -15421,10 +15493,14 @@ def test_update_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeaturestoreRequest.pb( featurestore_service.UpdateFeaturestoreRequest() ) @@ -15448,6 +15524,7 @@ def test_update_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_featurestore( request, @@ -15459,6 +15536,7 @@ def test_update_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_featurestore_rest_bad_request( @@ -15539,10 +15617,14 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeaturestoreRequest.pb( featurestore_service.DeleteFeaturestoreRequest() ) @@ -15566,6 +15648,7 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_featurestore( request, @@ -15577,6 +15660,7 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_entity_type_rest_bad_request( @@ -15749,10 +15833,14 @@ def test_create_entity_type_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateEntityTypeRequest.pb( featurestore_service.CreateEntityTypeRequest() ) @@ -15776,6 +15864,7 @@ def test_create_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_entity_type( request, @@ -15787,6 +15876,7 @@ def test_create_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_entity_type_rest_bad_request( @@ -15885,10 +15975,14 @@ def test_get_entity_type_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_get_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetEntityTypeRequest.pb( featurestore_service.GetEntityTypeRequest() ) @@ -15912,6 +16006,7 @@ def test_get_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = entity_type.EntityType() + post_with_metadata.return_value = entity_type.EntityType(), metadata client.get_entity_type( request, @@ -15923,6 +16018,7 @@ def test_get_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_entity_types_rest_bad_request( @@ -16011,10 +16107,14 @@ def test_list_entity_types_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_entity_types" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_entity_types_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_entity_types" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListEntityTypesRequest.pb( featurestore_service.ListEntityTypesRequest() ) @@ -16040,6 +16140,10 @@ def test_list_entity_types_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListEntityTypesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListEntityTypesResponse(), + metadata, + ) client.list_entity_types( request, @@ -16051,6 +16155,7 @@ def test_list_entity_types_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_entity_type_rest_bad_request( @@ -16241,10 +16346,14 @@ def test_update_entity_type_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateEntityTypeRequest.pb( featurestore_service.UpdateEntityTypeRequest() ) @@ -16268,6 +16377,7 @@ def test_update_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_entity_type.EntityType() + post_with_metadata.return_value = gca_entity_type.EntityType(), metadata client.update_entity_type( request, @@ -16279,6 +16389,7 @@ def test_update_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_entity_type_rest_bad_request( @@ -16363,10 +16474,14 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteEntityTypeRequest.pb( featurestore_service.DeleteEntityTypeRequest() ) @@ -16390,6 +16505,7 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_entity_type( request, @@ -16401,6 +16517,7 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_rest_bad_request( @@ -16578,10 +16695,14 @@ def test_create_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -16605,6 +16726,7 @@ def test_create_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature( request, @@ -16616,6 +16738,7 @@ def test_create_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_features_rest_bad_request( @@ -16700,10 +16823,14 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -16727,6 +16854,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_create_features( request, @@ -16738,6 +16866,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_rest_bad_request( @@ -16838,10 +16967,13 @@ def test_get_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, "post_get_feature_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -16865,6 +16997,7 @@ def test_get_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata client.get_feature( request, @@ -16876,6 +17009,7 @@ def test_get_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_features_rest_bad_request( @@ -16964,10 +17098,14 @@ def test_list_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -16993,6 +17131,10 @@ def test_list_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) client.list_features( request, @@ -17004,6 +17146,7 @@ def test_list_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_rest_bad_request( @@ -17201,10 +17344,14 @@ def test_update_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -17228,6 +17375,7 @@ def test_update_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_feature.Feature() + post_with_metadata.return_value = gca_feature.Feature(), metadata client.update_feature( request, @@ -17239,6 +17387,7 @@ def test_update_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_rest_bad_request( @@ -17323,10 +17472,14 @@ def test_delete_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -17350,6 +17503,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature( request, @@ -17361,6 +17515,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_feature_values_rest_bad_request( @@ -17445,10 +17600,14 @@ def test_import_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_import_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_import_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_import_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ImportFeatureValuesRequest.pb( featurestore_service.ImportFeatureValuesRequest() ) @@ -17472,6 +17631,7 @@ def test_import_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_feature_values( request, @@ -17483,6 +17643,7 @@ def test_import_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_read_feature_values_rest_bad_request( @@ -17567,10 +17728,14 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_batch_read_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_batch_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_batch_read_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchReadFeatureValuesRequest.pb( featurestore_service.BatchReadFeatureValuesRequest() ) @@ -17594,6 +17759,7 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_read_feature_values( request, @@ -17605,6 +17771,7 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_feature_values_rest_bad_request( @@ -17689,10 +17856,14 @@ def test_export_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_export_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_export_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_export_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ExportFeatureValuesRequest.pb( featurestore_service.ExportFeatureValuesRequest() ) @@ -17716,6 +17887,7 @@ def test_export_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_feature_values( request, @@ -17727,6 +17899,7 @@ def test_export_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_values_rest_bad_request( @@ -17811,10 +17984,14 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureValuesRequest.pb( featurestore_service.DeleteFeatureValuesRequest() ) @@ -17838,6 +18015,7 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_values( request, @@ -17849,6 +18027,7 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_features_rest_bad_request( @@ -17933,10 +18112,14 @@ def test_search_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_search_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_search_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_search_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.SearchFeaturesRequest.pb( featurestore_service.SearchFeaturesRequest() ) @@ -17962,6 +18145,10 @@ def test_search_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.SearchFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.SearchFeaturesResponse(), + metadata, + ) client.search_features( request, @@ -17973,6 +18160,7 @@ def test_search_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -19264,10 +19452,14 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeaturestoreRequest.pb( featurestore_service.CreateFeaturestoreRequest() ) @@ -19291,6 +19483,7 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_featurestore( request, @@ -19302,6 +19495,7 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19412,10 +19606,14 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeaturestoreRequest.pb( featurestore_service.GetFeaturestoreRequest() ) @@ -19439,6 +19637,7 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore.Featurestore() + post_with_metadata.return_value = featurestore.Featurestore(), metadata await client.get_featurestore( request, @@ -19450,6 +19649,7 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19550,10 +19750,14 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_featurestores" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_featurestores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_featurestores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturestoresRequest.pb( featurestore_service.ListFeaturestoresRequest() ) @@ -19579,6 +19783,10 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturestoresResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturestoresResponse(), + metadata, + ) await client.list_featurestores( request, @@ -19590,6 +19798,7 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19783,10 +19992,14 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeaturestoreRequest.pb( featurestore_service.UpdateFeaturestoreRequest() ) @@ -19810,6 +20023,7 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_featurestore( request, @@ -19821,6 +20035,7 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19917,10 +20132,14 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeaturestoreRequest.pb( featurestore_service.DeleteFeaturestoreRequest() ) @@ -19944,6 +20163,7 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_featurestore( request, @@ -19955,6 +20175,7 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20143,10 +20364,14 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateEntityTypeRequest.pb( featurestore_service.CreateEntityTypeRequest() ) @@ -20170,6 +20395,7 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_entity_type( request, @@ -20181,6 +20407,7 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20295,10 +20522,14 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetEntityTypeRequest.pb( featurestore_service.GetEntityTypeRequest() ) @@ -20322,6 +20553,7 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = entity_type.EntityType() + post_with_metadata.return_value = entity_type.EntityType(), metadata await client.get_entity_type( request, @@ -20333,6 +20565,7 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20437,10 +20670,14 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_entity_types" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_entity_types_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_entity_types" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListEntityTypesRequest.pb( featurestore_service.ListEntityTypesRequest() ) @@ -20466,6 +20703,10 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListEntityTypesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListEntityTypesResponse(), + metadata, + ) await client.list_entity_types( request, @@ -20477,6 +20718,7 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20683,10 +20925,14 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateEntityTypeRequest.pb( featurestore_service.UpdateEntityTypeRequest() ) @@ -20710,6 +20956,7 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_entity_type.EntityType() + post_with_metadata.return_value = gca_entity_type.EntityType(), metadata await client.update_entity_type( request, @@ -20721,6 +20968,7 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20821,10 +21069,14 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteEntityTypeRequest.pb( featurestore_service.DeleteEntityTypeRequest() ) @@ -20848,6 +21100,7 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_entity_type( request, @@ -20859,6 +21112,7 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21052,10 +21306,14 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -21079,6 +21337,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature( request, @@ -21090,6 +21349,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21190,10 +21450,14 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -21217,6 +21481,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_create_features( request, @@ -21228,6 +21493,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21344,10 +21610,14 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -21371,6 +21641,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata await client.get_feature( request, @@ -21382,6 +21653,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21486,10 +21758,14 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -21515,6 +21791,10 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) await client.list_features( request, @@ -21526,6 +21806,7 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21739,10 +22020,14 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -21766,6 +22051,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_feature.Feature() + post_with_metadata.return_value = gca_feature.Feature(), metadata await client.update_feature( request, @@ -21777,6 +22063,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21877,10 +22164,14 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -21904,6 +22195,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature( request, @@ -21915,6 +22207,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22015,10 +22308,14 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_import_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_import_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_import_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ImportFeatureValuesRequest.pb( featurestore_service.ImportFeatureValuesRequest() ) @@ -22042,6 +22339,7 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_feature_values( request, @@ -22053,6 +22351,7 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22154,11 +22453,15 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep transports.AsyncFeaturestoreServiceRestInterceptor, "post_batch_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_batch_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_batch_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchReadFeatureValuesRequest.pb( featurestore_service.BatchReadFeatureValuesRequest() ) @@ -22182,6 +22485,7 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_read_feature_values( request, @@ -22193,6 +22497,7 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22293,10 +22598,14 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_export_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_export_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_export_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ExportFeatureValuesRequest.pb( featurestore_service.ExportFeatureValuesRequest() ) @@ -22320,6 +22629,7 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_feature_values( request, @@ -22331,6 +22641,7 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22431,10 +22742,14 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureValuesRequest.pb( featurestore_service.DeleteFeatureValuesRequest() ) @@ -22458,6 +22773,7 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_values( request, @@ -22469,6 +22785,7 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22569,10 +22886,14 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_search_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_search_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_search_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.SearchFeaturesRequest.pb( featurestore_service.SearchFeaturesRequest() ) @@ -22598,6 +22919,10 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.SearchFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.SearchFeaturesResponse(), + metadata, + ) await client.search_features( request, @@ -22609,6 +22934,7 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_gen_ai_cache_service.py b/tests/unit/gapic/aiplatform_v1/test_gen_ai_cache_service.py index efbe3252a7..8618021b79 100644 --- a/tests/unit/gapic/aiplatform_v1/test_gen_ai_cache_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_gen_ai_cache_service.py @@ -90,6 +90,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -356,6 +364,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = GenAiCacheServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = GenAiCacheServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4788,10 +4839,14 @@ def test_create_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_create_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_create_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_create_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.CreateCachedContentRequest.pb( gen_ai_cache_service.CreateCachedContentRequest() ) @@ -4817,6 +4872,7 @@ def test_create_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata client.create_cached_content( request, @@ -4828,6 +4884,7 @@ def test_create_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_cached_content_rest_bad_request( @@ -4916,10 +4973,14 @@ def test_get_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_get_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_get_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_get_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.GetCachedContentRequest.pb( gen_ai_cache_service.GetCachedContentRequest() ) @@ -4945,6 +5006,7 @@ def test_get_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = cached_content.CachedContent() + post_with_metadata.return_value = cached_content.CachedContent(), metadata client.get_cached_content( request, @@ -4956,6 +5018,7 @@ def test_get_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_cached_content_rest_bad_request( @@ -5246,10 +5309,14 @@ def test_update_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_update_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_update_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_update_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.UpdateCachedContentRequest.pb( gen_ai_cache_service.UpdateCachedContentRequest() ) @@ -5275,6 +5342,7 @@ def test_update_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata client.update_cached_content( request, @@ -5286,6 +5354,7 @@ def test_update_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_cached_content_rest_bad_request( @@ -5479,10 +5548,14 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_list_cached_contents" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_list_cached_contents_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_list_cached_contents" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.ListCachedContentsRequest.pb( gen_ai_cache_service.ListCachedContentsRequest() ) @@ -5508,6 +5581,10 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gen_ai_cache_service.ListCachedContentsResponse() + post_with_metadata.return_value = ( + gen_ai_cache_service.ListCachedContentsResponse(), + metadata, + ) client.list_cached_contents( request, @@ -5519,6 +5596,7 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6570,10 +6648,14 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_create_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_create_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_create_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.CreateCachedContentRequest.pb( gen_ai_cache_service.CreateCachedContentRequest() ) @@ -6599,6 +6681,7 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata await client.create_cached_content( request, @@ -6610,6 +6693,7 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6714,10 +6798,14 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_get_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_get_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_get_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.GetCachedContentRequest.pb( gen_ai_cache_service.GetCachedContentRequest() ) @@ -6743,6 +6831,7 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = cached_content.CachedContent() + post_with_metadata.return_value = cached_content.CachedContent(), metadata await client.get_cached_content( request, @@ -6754,6 +6843,7 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7060,10 +7150,14 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_update_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_update_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_update_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.UpdateCachedContentRequest.pb( gen_ai_cache_service.UpdateCachedContentRequest() ) @@ -7089,6 +7183,7 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata await client.update_cached_content( request, @@ -7100,6 +7195,7 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7325,10 +7421,14 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_list_cached_contents" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_list_cached_contents_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_list_cached_contents" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.ListCachedContentsRequest.pb( gen_ai_cache_service.ListCachedContentsRequest() ) @@ -7354,6 +7454,10 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gen_ai_cache_service.ListCachedContentsResponse() + post_with_metadata.return_value = ( + gen_ai_cache_service.ListCachedContentsResponse(), + metadata, + ) await client.list_cached_contents( request, @@ -7365,6 +7469,7 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_gen_ai_tuning_service.py b/tests/unit/gapic/aiplatform_v1/test_gen_ai_tuning_service.py index d759daf0ac..69c3bf3309 100644 --- a/tests/unit/gapic/aiplatform_v1/test_gen_ai_tuning_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_gen_ai_tuning_service.py @@ -96,6 +96,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -366,6 +374,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = GenAiTuningServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = GenAiTuningServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4699,10 +4750,14 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_create_tuning_job" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_create_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_create_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.CreateTuningJobRequest.pb( genai_tuning_service.CreateTuningJobRequest() ) @@ -4726,6 +4781,7 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tuning_job.TuningJob() + post_with_metadata.return_value = gca_tuning_job.TuningJob(), metadata client.create_tuning_job( request, @@ -4737,6 +4793,7 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tuning_job_rest_bad_request( @@ -4832,10 +4889,14 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_get_tuning_job" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_get_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_get_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.GetTuningJobRequest.pb( genai_tuning_service.GetTuningJobRequest() ) @@ -4859,6 +4920,7 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tuning_job.TuningJob() + post_with_metadata.return_value = tuning_job.TuningJob(), metadata client.get_tuning_job( request, @@ -4870,6 +4932,7 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tuning_jobs_rest_bad_request( @@ -4954,10 +5017,14 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_list_tuning_jobs" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_list_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_list_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.ListTuningJobsRequest.pb( genai_tuning_service.ListTuningJobsRequest() ) @@ -4983,6 +5050,10 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = genai_tuning_service.ListTuningJobsResponse() + post_with_metadata.return_value = ( + genai_tuning_service.ListTuningJobsResponse(), + metadata, + ) client.list_tuning_jobs( request, @@ -4994,6 +5065,7 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_tuning_job_rest_bad_request( @@ -5183,10 +5255,14 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_rebase_tuned_model" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_rebase_tuned_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_rebase_tuned_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.RebaseTunedModelRequest.pb( genai_tuning_service.RebaseTunedModelRequest() ) @@ -5210,6 +5286,7 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.rebase_tuned_model( request, @@ -5221,6 +5298,7 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6262,10 +6340,14 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_create_tuning_job" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_create_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_create_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.CreateTuningJobRequest.pb( genai_tuning_service.CreateTuningJobRequest() ) @@ -6289,6 +6371,7 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tuning_job.TuningJob() + post_with_metadata.return_value = gca_tuning_job.TuningJob(), metadata await client.create_tuning_job( request, @@ -6300,6 +6383,7 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6411,10 +6495,14 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_get_tuning_job" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_get_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_get_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.GetTuningJobRequest.pb( genai_tuning_service.GetTuningJobRequest() ) @@ -6438,6 +6526,7 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tuning_job.TuningJob() + post_with_metadata.return_value = tuning_job.TuningJob(), metadata await client.get_tuning_job( request, @@ -6449,6 +6538,7 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6549,10 +6639,14 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_list_tuning_jobs" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_list_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_list_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.ListTuningJobsRequest.pb( genai_tuning_service.ListTuningJobsRequest() ) @@ -6578,6 +6672,10 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = genai_tuning_service.ListTuningJobsResponse() + post_with_metadata.return_value = ( + genai_tuning_service.ListTuningJobsResponse(), + metadata, + ) await client.list_tuning_jobs( request, @@ -6589,6 +6687,7 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6810,10 +6909,14 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_rebase_tuned_model" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_rebase_tuned_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_rebase_tuned_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.RebaseTunedModelRequest.pb( genai_tuning_service.RebaseTunedModelRequest() ) @@ -6837,6 +6940,7 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.rebase_tuned_model( request, @@ -6848,6 +6952,7 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_index_endpoint_service.py b/tests/unit/gapic/aiplatform_v1/test_index_endpoint_service.py index 7375afc246..0a5dc7653d 100644 --- a/tests/unit/gapic/aiplatform_v1/test_index_endpoint_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_index_endpoint_service.py @@ -95,6 +95,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -367,6 +375,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = IndexEndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = IndexEndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -6584,10 +6635,14 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_create_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_create_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_create_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.CreateIndexEndpointRequest.pb( index_endpoint_service.CreateIndexEndpointRequest() ) @@ -6611,6 +6666,7 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_index_endpoint( request, @@ -6622,6 +6678,7 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_index_endpoint_rest_bad_request( @@ -6724,10 +6781,14 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_get_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_get_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_get_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.GetIndexEndpointRequest.pb( index_endpoint_service.GetIndexEndpointRequest() ) @@ -6753,6 +6814,7 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint.IndexEndpoint() + post_with_metadata.return_value = index_endpoint.IndexEndpoint(), metadata client.get_index_endpoint( request, @@ -6764,6 +6826,7 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_index_endpoints_rest_bad_request( @@ -6850,10 +6913,14 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_list_index_endpoints" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_list_index_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_list_index_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.ListIndexEndpointsRequest.pb( index_endpoint_service.ListIndexEndpointsRequest() ) @@ -6879,6 +6946,10 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint_service.ListIndexEndpointsResponse() + post_with_metadata.return_value = ( + index_endpoint_service.ListIndexEndpointsResponse(), + metadata, + ) client.list_index_endpoints( request, @@ -6890,6 +6961,7 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_index_endpoint_rest_bad_request( @@ -7155,10 +7227,14 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_update_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_update_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_update_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UpdateIndexEndpointRequest.pb( index_endpoint_service.UpdateIndexEndpointRequest() ) @@ -7184,6 +7260,7 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_index_endpoint.IndexEndpoint() + post_with_metadata.return_value = gca_index_endpoint.IndexEndpoint(), metadata client.update_index_endpoint( request, @@ -7195,6 +7272,7 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_index_endpoint_rest_bad_request( @@ -7275,10 +7353,14 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_delete_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_delete_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_delete_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeleteIndexEndpointRequest.pb( index_endpoint_service.DeleteIndexEndpointRequest() ) @@ -7302,6 +7384,7 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_index_endpoint( request, @@ -7313,6 +7396,7 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_deploy_index_rest_bad_request( @@ -7397,10 +7481,14 @@ def test_deploy_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_deploy_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_deploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_deploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeployIndexRequest.pb( index_endpoint_service.DeployIndexRequest() ) @@ -7424,6 +7512,7 @@ def test_deploy_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.deploy_index( request, @@ -7435,6 +7524,7 @@ def test_deploy_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_undeploy_index_rest_bad_request( @@ -7519,10 +7609,14 @@ def test_undeploy_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_undeploy_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_undeploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_undeploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UndeployIndexRequest.pb( index_endpoint_service.UndeployIndexRequest() ) @@ -7546,6 +7640,7 @@ def test_undeploy_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.undeploy_index( request, @@ -7557,6 +7652,7 @@ def test_undeploy_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_mutate_deployed_index_rest_bad_request( @@ -7764,10 +7860,14 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_mutate_deployed_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_mutate_deployed_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_mutate_deployed_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.MutateDeployedIndexRequest.pb( index_endpoint_service.MutateDeployedIndexRequest() ) @@ -7791,6 +7891,7 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.mutate_deployed_index( request, @@ -7802,6 +7903,7 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -8886,10 +8988,14 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_create_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_create_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_create_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.CreateIndexEndpointRequest.pb( index_endpoint_service.CreateIndexEndpointRequest() ) @@ -8913,6 +9019,7 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_index_endpoint( request, @@ -8924,6 +9031,7 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9042,10 +9150,14 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_get_index_endpoint" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_get_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_get_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.GetIndexEndpointRequest.pb( index_endpoint_service.GetIndexEndpointRequest() ) @@ -9071,6 +9183,7 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint.IndexEndpoint() + post_with_metadata.return_value = index_endpoint.IndexEndpoint(), metadata await client.get_index_endpoint( request, @@ -9082,6 +9195,7 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9184,10 +9298,14 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_list_index_endpoints" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_list_index_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_list_index_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.ListIndexEndpointsRequest.pb( index_endpoint_service.ListIndexEndpointsRequest() ) @@ -9213,6 +9331,10 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint_service.ListIndexEndpointsResponse() + post_with_metadata.return_value = ( + index_endpoint_service.ListIndexEndpointsResponse(), + metadata, + ) await client.list_index_endpoints( request, @@ -9224,6 +9346,7 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9506,10 +9629,14 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_update_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_update_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_update_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UpdateIndexEndpointRequest.pb( index_endpoint_service.UpdateIndexEndpointRequest() ) @@ -9535,6 +9662,7 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_index_endpoint.IndexEndpoint() + post_with_metadata.return_value = gca_index_endpoint.IndexEndpoint(), metadata await client.update_index_endpoint( request, @@ -9546,6 +9674,7 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9643,10 +9772,14 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_delete_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_delete_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_delete_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeleteIndexEndpointRequest.pb( index_endpoint_service.DeleteIndexEndpointRequest() ) @@ -9670,6 +9803,7 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_index_endpoint( request, @@ -9681,6 +9815,7 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9781,10 +9916,14 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_deploy_index" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_deploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_deploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeployIndexRequest.pb( index_endpoint_service.DeployIndexRequest() ) @@ -9808,6 +9947,7 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.deploy_index( request, @@ -9819,6 +9959,7 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9919,10 +10060,14 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_undeploy_index" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_undeploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_undeploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UndeployIndexRequest.pb( index_endpoint_service.UndeployIndexRequest() ) @@ -9946,6 +10091,7 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.undeploy_index( request, @@ -9957,6 +10103,7 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10181,10 +10328,14 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_mutate_deployed_index", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_mutate_deployed_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_mutate_deployed_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.MutateDeployedIndexRequest.pb( index_endpoint_service.MutateDeployedIndexRequest() ) @@ -10208,6 +10359,7 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.mutate_deployed_index( request, @@ -10219,6 +10371,7 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_index_service.py b/tests/unit/gapic/aiplatform_v1/test_index_service.py index 7235f715ed..30892adb24 100644 --- a/tests/unit/gapic/aiplatform_v1/test_index_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_index_service.py @@ -89,6 +89,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -331,6 +339,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = IndexServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = IndexServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5287,10 +5338,13 @@ def test_create_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_create_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_create_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_create_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.CreateIndexRequest.pb( index_service.CreateIndexRequest() ) @@ -5314,6 +5368,7 @@ def test_create_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_index( request, @@ -5325,6 +5380,7 @@ def test_create_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_index_rest_bad_request(request_type=index_service.GetIndexRequest): @@ -5421,10 +5477,13 @@ def test_get_index_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_get_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_get_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_get_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.GetIndexRequest.pb(index_service.GetIndexRequest()) transcode.return_value = { "method": "post", @@ -5446,6 +5505,7 @@ def test_get_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index.Index() + post_with_metadata.return_value = index.Index(), metadata client.get_index( request, @@ -5457,6 +5517,7 @@ def test_get_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_indexes_rest_bad_request(request_type=index_service.ListIndexesRequest): @@ -5539,10 +5600,13 @@ def test_list_indexes_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_list_indexes" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_list_indexes_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_list_indexes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.ListIndexesRequest.pb( index_service.ListIndexesRequest() ) @@ -5568,6 +5632,7 @@ def test_list_indexes_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.ListIndexesResponse() + post_with_metadata.return_value = index_service.ListIndexesResponse(), metadata client.list_indexes( request, @@ -5579,6 +5644,7 @@ def test_list_indexes_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_index_rest_bad_request(request_type=index_service.UpdateIndexRequest): @@ -5762,10 +5828,13 @@ def test_update_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_update_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_update_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_update_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpdateIndexRequest.pb( index_service.UpdateIndexRequest() ) @@ -5789,6 +5858,7 @@ def test_update_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_index( request, @@ -5800,6 +5870,7 @@ def test_update_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_index_rest_bad_request(request_type=index_service.DeleteIndexRequest): @@ -5878,10 +5949,13 @@ def test_delete_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_delete_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_delete_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_delete_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.DeleteIndexRequest.pb( index_service.DeleteIndexRequest() ) @@ -5905,6 +5979,7 @@ def test_delete_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_index( request, @@ -5916,6 +5991,7 @@ def test_delete_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upsert_datapoints_rest_bad_request( @@ -5997,10 +6073,13 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_upsert_datapoints" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_upsert_datapoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_upsert_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpsertDatapointsRequest.pb( index_service.UpsertDatapointsRequest() ) @@ -6026,6 +6105,10 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.UpsertDatapointsResponse() + post_with_metadata.return_value = ( + index_service.UpsertDatapointsResponse(), + metadata, + ) client.upsert_datapoints( request, @@ -6037,6 +6120,7 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_remove_datapoints_rest_bad_request( @@ -6118,10 +6202,13 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_remove_datapoints" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_remove_datapoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_remove_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.RemoveDatapointsRequest.pb( index_service.RemoveDatapointsRequest() ) @@ -6147,6 +6234,10 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.RemoveDatapointsResponse() + post_with_metadata.return_value = ( + index_service.RemoveDatapointsResponse(), + metadata, + ) client.remove_datapoints( request, @@ -6158,6 +6249,7 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7159,10 +7251,13 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_create_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_create_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_create_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.CreateIndexRequest.pb( index_service.CreateIndexRequest() ) @@ -7186,6 +7281,7 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_index( request, @@ -7197,6 +7293,7 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7311,10 +7408,13 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_get_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_get_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_get_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.GetIndexRequest.pb(index_service.GetIndexRequest()) transcode.return_value = { "method": "post", @@ -7336,6 +7436,7 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index.Index() + post_with_metadata.return_value = index.Index(), metadata await client.get_index( request, @@ -7347,6 +7448,7 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7447,10 +7549,13 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_list_indexes" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_list_indexes_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_list_indexes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.ListIndexesRequest.pb( index_service.ListIndexesRequest() ) @@ -7476,6 +7581,7 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.ListIndexesResponse() + post_with_metadata.return_value = index_service.ListIndexesResponse(), metadata await client.list_indexes( request, @@ -7487,6 +7593,7 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7688,10 +7795,13 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_update_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_update_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_update_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpdateIndexRequest.pb( index_service.UpdateIndexRequest() ) @@ -7715,6 +7825,7 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_index( request, @@ -7726,6 +7837,7 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7822,10 +7934,13 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_delete_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_delete_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_delete_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.DeleteIndexRequest.pb( index_service.DeleteIndexRequest() ) @@ -7849,6 +7964,7 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_index( request, @@ -7860,6 +7976,7 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7957,10 +8074,14 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_upsert_datapoints" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, + "post_upsert_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_upsert_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpsertDatapointsRequest.pb( index_service.UpsertDatapointsRequest() ) @@ -7986,6 +8107,10 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.UpsertDatapointsResponse() + post_with_metadata.return_value = ( + index_service.UpsertDatapointsResponse(), + metadata, + ) await client.upsert_datapoints( request, @@ -7997,6 +8122,7 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8094,10 +8220,14 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_remove_datapoints" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, + "post_remove_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_remove_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.RemoveDatapointsRequest.pb( index_service.RemoveDatapointsRequest() ) @@ -8123,6 +8253,10 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.RemoveDatapointsResponse() + post_with_metadata.return_value = ( + index_service.RemoveDatapointsResponse(), + metadata, + ) await client.remove_datapoints( request, @@ -8134,6 +8268,7 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_job_service.py b/tests/unit/gapic/aiplatform_v1/test_job_service.py index cb9526aaaa..e8f8b7df8a 100644 --- a/tests/unit/gapic/aiplatform_v1/test_job_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_job_service.py @@ -123,6 +123,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -357,6 +365,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = JobServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = JobServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -24530,10 +24581,13 @@ def test_create_custom_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_create_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateCustomJobRequest.pb( job_service.CreateCustomJobRequest() ) @@ -24557,6 +24611,7 @@ def test_create_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_custom_job.CustomJob() + post_with_metadata.return_value = gca_custom_job.CustomJob(), metadata client.create_custom_job( request, @@ -24568,6 +24623,7 @@ def test_create_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_custom_job_rest_bad_request(request_type=job_service.GetCustomJobRequest): @@ -24658,10 +24714,13 @@ def test_get_custom_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetCustomJobRequest.pb( job_service.GetCustomJobRequest() ) @@ -24685,6 +24744,7 @@ def test_get_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = custom_job.CustomJob() + post_with_metadata.return_value = custom_job.CustomJob(), metadata client.get_custom_job( request, @@ -24696,6 +24756,7 @@ def test_get_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_custom_jobs_rest_bad_request( @@ -24780,10 +24841,13 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_custom_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_list_custom_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_custom_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListCustomJobsRequest.pb( job_service.ListCustomJobsRequest() ) @@ -24809,6 +24873,7 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListCustomJobsResponse() + post_with_metadata.return_value = job_service.ListCustomJobsResponse(), metadata client.list_custom_jobs( request, @@ -24820,6 +24885,7 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_custom_job_rest_bad_request( @@ -24900,10 +24966,13 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_delete_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteCustomJobRequest.pb( job_service.DeleteCustomJobRequest() ) @@ -24927,6 +24996,7 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_custom_job( request, @@ -24938,6 +25008,7 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_custom_job_rest_bad_request( @@ -25265,10 +25336,14 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateDataLabelingJobRequest.pb( job_service.CreateDataLabelingJobRequest() ) @@ -25294,6 +25369,10 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = ( + gca_data_labeling_job.DataLabelingJob(), + metadata, + ) client.create_data_labeling_job( request, @@ -25305,6 +25384,7 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_data_labeling_job_rest_bad_request( @@ -25409,10 +25489,13 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_data_labeling_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetDataLabelingJobRequest.pb( job_service.GetDataLabelingJobRequest() ) @@ -25438,6 +25521,7 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = data_labeling_job.DataLabelingJob(), metadata client.get_data_labeling_job( request, @@ -25449,6 +25533,7 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_data_labeling_jobs_rest_bad_request( @@ -25533,10 +25618,14 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_data_labeling_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_data_labeling_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_data_labeling_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListDataLabelingJobsRequest.pb( job_service.ListDataLabelingJobsRequest() ) @@ -25562,6 +25651,10 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListDataLabelingJobsResponse() + post_with_metadata.return_value = ( + job_service.ListDataLabelingJobsResponse(), + metadata, + ) client.list_data_labeling_jobs( request, @@ -25573,6 +25666,7 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_data_labeling_job_rest_bad_request( @@ -25657,10 +25751,14 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteDataLabelingJobRequest.pb( job_service.DeleteDataLabelingJobRequest() ) @@ -25684,6 +25782,7 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_data_labeling_job( request, @@ -25695,6 +25794,7 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_data_labeling_job_rest_bad_request( @@ -26174,10 +26274,14 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateHyperparameterTuningJobRequest.pb( job_service.CreateHyperparameterTuningJobRequest() ) @@ -26203,6 +26307,10 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + gca_hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) client.create_hyperparameter_tuning_job( request, @@ -26214,6 +26322,7 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_hyperparameter_tuning_job_rest_bad_request( @@ -26318,10 +26427,14 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetHyperparameterTuningJobRequest.pb( job_service.GetHyperparameterTuningJobRequest() ) @@ -26347,6 +26460,10 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) client.get_hyperparameter_tuning_job( request, @@ -26358,6 +26475,7 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_hyperparameter_tuning_jobs_rest_bad_request( @@ -26442,10 +26560,14 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_hyperparameter_tuning_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_hyperparameter_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_hyperparameter_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListHyperparameterTuningJobsRequest.pb( job_service.ListHyperparameterTuningJobsRequest() ) @@ -26471,6 +26593,10 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListHyperparameterTuningJobsResponse() + post_with_metadata.return_value = ( + job_service.ListHyperparameterTuningJobsResponse(), + metadata, + ) client.list_hyperparameter_tuning_jobs( request, @@ -26482,6 +26608,7 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_hyperparameter_tuning_job_rest_bad_request( @@ -26566,10 +26693,14 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteHyperparameterTuningJobRequest.pb( job_service.DeleteHyperparameterTuningJobRequest() ) @@ -26593,6 +26724,7 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_hyperparameter_tuning_job( request, @@ -26604,6 +26736,7 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_hyperparameter_tuning_job_rest_bad_request( @@ -27009,10 +27142,13 @@ def test_create_nas_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_create_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateNasJobRequest.pb( job_service.CreateNasJobRequest() ) @@ -27036,6 +27172,7 @@ def test_create_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_nas_job.NasJob() + post_with_metadata.return_value = gca_nas_job.NasJob(), metadata client.create_nas_job( request, @@ -27047,6 +27184,7 @@ def test_create_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_nas_job_rest_bad_request(request_type=job_service.GetNasJobRequest): @@ -27139,10 +27277,13 @@ def test_get_nas_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasJobRequest.pb(job_service.GetNasJobRequest()) transcode.return_value = { "method": "post", @@ -27164,6 +27305,7 @@ def test_get_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasJob() + post_with_metadata.return_value = nas_job.NasJob(), metadata client.get_nas_job( request, @@ -27175,6 +27317,7 @@ def test_get_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_nas_jobs_rest_bad_request(request_type=job_service.ListNasJobsRequest): @@ -27257,10 +27400,13 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_nas_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_list_nas_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_nas_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasJobsRequest.pb(job_service.ListNasJobsRequest()) transcode.return_value = { "method": "post", @@ -27284,6 +27430,7 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasJobsResponse() + post_with_metadata.return_value = job_service.ListNasJobsResponse(), metadata client.list_nas_jobs( request, @@ -27295,6 +27442,7 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_nas_job_rest_bad_request(request_type=job_service.DeleteNasJobRequest): @@ -27373,10 +27521,13 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_delete_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteNasJobRequest.pb( job_service.DeleteNasJobRequest() ) @@ -27400,6 +27551,7 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_nas_job( request, @@ -27411,6 +27563,7 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_nas_job_rest_bad_request(request_type=job_service.CancelNasJobRequest): @@ -27608,10 +27761,13 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_nas_trial_detail" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_nas_trial_detail_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_nas_trial_detail" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasTrialDetailRequest.pb( job_service.GetNasTrialDetailRequest() ) @@ -27635,6 +27791,7 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasTrialDetail() + post_with_metadata.return_value = nas_job.NasTrialDetail(), metadata client.get_nas_trial_detail( request, @@ -27646,6 +27803,7 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_nas_trial_details_rest_bad_request( @@ -27730,10 +27888,14 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_nas_trial_details" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_nas_trial_details_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_nas_trial_details" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasTrialDetailsRequest.pb( job_service.ListNasTrialDetailsRequest() ) @@ -27759,6 +27921,10 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasTrialDetailsResponse() + post_with_metadata.return_value = ( + job_service.ListNasTrialDetailsResponse(), + metadata, + ) client.list_nas_trial_details( request, @@ -27770,6 +27936,7 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_batch_prediction_job_rest_bad_request( @@ -27836,10 +28003,25 @@ def test_create_batch_prediction_job_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, }, "input_config": { @@ -28085,10 +28267,14 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateBatchPredictionJobRequest.pb( job_service.CreateBatchPredictionJobRequest() ) @@ -28114,6 +28300,10 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + gca_batch_prediction_job.BatchPredictionJob(), + metadata, + ) client.create_batch_prediction_job( request, @@ -28125,6 +28315,7 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_batch_prediction_job_rest_bad_request( @@ -28231,10 +28422,14 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetBatchPredictionJobRequest.pb( job_service.GetBatchPredictionJobRequest() ) @@ -28260,6 +28455,10 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + batch_prediction_job.BatchPredictionJob(), + metadata, + ) client.get_batch_prediction_job( request, @@ -28271,6 +28470,7 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_batch_prediction_jobs_rest_bad_request( @@ -28355,10 +28555,14 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_batch_prediction_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_batch_prediction_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_batch_prediction_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListBatchPredictionJobsRequest.pb( job_service.ListBatchPredictionJobsRequest() ) @@ -28384,6 +28588,10 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListBatchPredictionJobsResponse() + post_with_metadata.return_value = ( + job_service.ListBatchPredictionJobsResponse(), + metadata, + ) client.list_batch_prediction_jobs( request, @@ -28395,6 +28603,7 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_batch_prediction_job_rest_bad_request( @@ -28479,10 +28688,14 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteBatchPredictionJobRequest.pb( job_service.DeleteBatchPredictionJobRequest() ) @@ -28506,6 +28719,7 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_batch_prediction_job( request, @@ -28517,6 +28731,7 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_batch_prediction_job_rest_bad_request( @@ -28916,11 +29131,15 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_create_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateModelDeploymentMonitoringJobRequest.pb( job_service.CreateModelDeploymentMonitoringJobRequest() ) @@ -28950,6 +29169,10 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept post.return_value = ( gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) client.create_model_deployment_monitoring_job( request, @@ -28961,6 +29184,7 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_model_deployment_monitoring_stats_anomalies_rest_bad_request( @@ -29062,11 +29286,15 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( transports.JobServiceRestInterceptor, "post_search_model_deployment_monitoring_stats_anomalies", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_search_model_deployment_monitoring_stats_anomalies_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_search_model_deployment_monitoring_stats_anomalies", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest.pb( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest() @@ -29098,6 +29326,10 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( post.return_value = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse() ) + post_with_metadata.return_value = ( + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse(), + metadata, + ) client.search_model_deployment_monitoring_stats_anomalies( request, @@ -29109,6 +29341,7 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_deployment_monitoring_job_rest_bad_request( @@ -29222,10 +29455,14 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_model_deployment_monitoring_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_model_deployment_monitoring_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetModelDeploymentMonitoringJobRequest.pb( job_service.GetModelDeploymentMonitoringJobRequest() ) @@ -29255,6 +29492,10 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) post.return_value = ( model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) client.get_model_deployment_monitoring_job( request, @@ -29266,6 +29507,7 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_deployment_monitoring_jobs_rest_bad_request( @@ -29353,11 +29595,15 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto transports.JobServiceRestInterceptor, "post_list_model_deployment_monitoring_jobs", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_model_deployment_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_model_deployment_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListModelDeploymentMonitoringJobsRequest.pb( job_service.ListModelDeploymentMonitoringJobsRequest() ) @@ -29383,6 +29629,10 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = job_service.ListModelDeploymentMonitoringJobsResponse() + post_with_metadata.return_value = ( + job_service.ListModelDeploymentMonitoringJobsResponse(), + metadata, + ) client.list_model_deployment_monitoring_jobs( request, @@ -29394,6 +29644,7 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_model_deployment_monitoring_job_rest_bad_request( @@ -29657,11 +29908,15 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_update_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_update_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_update_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.UpdateModelDeploymentMonitoringJobRequest.pb( job_service.UpdateModelDeploymentMonitoringJobRequest() ) @@ -29685,6 +29940,7 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_model_deployment_monitoring_job( request, @@ -29696,6 +29952,7 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_deployment_monitoring_job_rest_bad_request( @@ -29781,11 +30038,15 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_delete_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteModelDeploymentMonitoringJobRequest.pb( job_service.DeleteModelDeploymentMonitoringJobRequest() ) @@ -29809,6 +30070,7 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_deployment_monitoring_job( request, @@ -29820,6 +30082,7 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_pause_model_deployment_monitoring_job_rest_bad_request( @@ -31728,10 +31991,14 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_custom_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateCustomJobRequest.pb( job_service.CreateCustomJobRequest() ) @@ -31755,6 +32022,7 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_custom_job.CustomJob() + post_with_metadata.return_value = gca_custom_job.CustomJob(), metadata await client.create_custom_job( request, @@ -31766,6 +32034,7 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31874,10 +32143,13 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_get_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetCustomJobRequest.pb( job_service.GetCustomJobRequest() ) @@ -31901,6 +32173,7 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = custom_job.CustomJob() + post_with_metadata.return_value = custom_job.CustomJob(), metadata await client.get_custom_job( request, @@ -31912,6 +32185,7 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32012,10 +32286,13 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_custom_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_list_custom_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_custom_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListCustomJobsRequest.pb( job_service.ListCustomJobsRequest() ) @@ -32041,6 +32318,7 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListCustomJobsResponse() + post_with_metadata.return_value = job_service.ListCustomJobsResponse(), metadata await client.list_custom_jobs( request, @@ -32052,6 +32330,7 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32148,10 +32427,14 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_custom_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteCustomJobRequest.pb( job_service.DeleteCustomJobRequest() ) @@ -32175,6 +32458,7 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_custom_job( request, @@ -32186,6 +32470,7 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32545,10 +32830,14 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateDataLabelingJobRequest.pb( job_service.CreateDataLabelingJobRequest() ) @@ -32574,6 +32863,10 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = gca_data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = ( + gca_data_labeling_job.DataLabelingJob(), + metadata, + ) await client.create_data_labeling_job( request, @@ -32585,6 +32878,7 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32705,10 +32999,14 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetDataLabelingJobRequest.pb( job_service.GetDataLabelingJobRequest() ) @@ -32734,6 +33032,7 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = data_labeling_job.DataLabelingJob(), metadata await client.get_data_labeling_job( request, @@ -32745,6 +33044,7 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32845,10 +33145,14 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_data_labeling_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_data_labeling_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_data_labeling_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListDataLabelingJobsRequest.pb( job_service.ListDataLabelingJobsRequest() ) @@ -32874,6 +33178,10 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = job_service.ListDataLabelingJobsResponse() + post_with_metadata.return_value = ( + job_service.ListDataLabelingJobsResponse(), + metadata, + ) await client.list_data_labeling_jobs( request, @@ -32885,6 +33193,7 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32985,10 +33294,14 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteDataLabelingJobRequest.pb( job_service.DeleteDataLabelingJobRequest() ) @@ -33012,6 +33325,7 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_data_labeling_job( request, @@ -33023,6 +33337,7 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33537,11 +33852,15 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_create_hyperparameter_tuning_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_hyperparameter_tuning_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateHyperparameterTuningJobRequest.pb( job_service.CreateHyperparameterTuningJobRequest() ) @@ -33567,6 +33886,10 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + gca_hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) await client.create_hyperparameter_tuning_job( request, @@ -33578,6 +33901,7 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33700,10 +34024,14 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetHyperparameterTuningJobRequest.pb( job_service.GetHyperparameterTuningJobRequest() ) @@ -33729,6 +34057,10 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) await client.get_hyperparameter_tuning_job( request, @@ -33740,6 +34072,7 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33843,10 +34176,14 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_list_hyperparameter_tuning_jobs", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_hyperparameter_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_hyperparameter_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListHyperparameterTuningJobsRequest.pb( job_service.ListHyperparameterTuningJobsRequest() ) @@ -33872,6 +34209,10 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = job_service.ListHyperparameterTuningJobsResponse() + post_with_metadata.return_value = ( + job_service.ListHyperparameterTuningJobsResponse(), + metadata, + ) await client.list_hyperparameter_tuning_jobs( request, @@ -33883,6 +34224,7 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33986,11 +34328,15 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_delete_hyperparameter_tuning_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_hyperparameter_tuning_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteHyperparameterTuningJobRequest.pb( job_service.DeleteHyperparameterTuningJobRequest() ) @@ -34014,6 +34360,7 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_hyperparameter_tuning_job( request, @@ -34025,6 +34372,7 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34467,10 +34815,13 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_create_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateNasJobRequest.pb( job_service.CreateNasJobRequest() ) @@ -34494,6 +34845,7 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_nas_job.NasJob() + post_with_metadata.return_value = gca_nas_job.NasJob(), metadata await client.create_nas_job( request, @@ -34505,6 +34857,7 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34615,10 +34968,13 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_get_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasJobRequest.pb(job_service.GetNasJobRequest()) transcode.return_value = { "method": "post", @@ -34640,6 +34996,7 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasJob() + post_with_metadata.return_value = nas_job.NasJob(), metadata await client.get_nas_job( request, @@ -34651,6 +35008,7 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34751,10 +35109,13 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_nas_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_list_nas_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_nas_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasJobsRequest.pb(job_service.ListNasJobsRequest()) transcode.return_value = { "method": "post", @@ -34778,6 +35139,7 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasJobsResponse() + post_with_metadata.return_value = job_service.ListNasJobsResponse(), metadata await client.list_nas_jobs( request, @@ -34789,6 +35151,7 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34885,10 +35248,13 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_delete_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteNasJobRequest.pb( job_service.DeleteNasJobRequest() ) @@ -34912,6 +35278,7 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_nas_job( request, @@ -34923,6 +35290,7 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35154,10 +35522,14 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_nas_trial_detail" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_nas_trial_detail_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_nas_trial_detail" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasTrialDetailRequest.pb( job_service.GetNasTrialDetailRequest() ) @@ -35181,6 +35553,7 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasTrialDetail() + post_with_metadata.return_value = nas_job.NasTrialDetail(), metadata await client.get_nas_trial_detail( request, @@ -35192,6 +35565,7 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35292,10 +35666,14 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_nas_trial_details" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_nas_trial_details_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_nas_trial_details" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasTrialDetailsRequest.pb( job_service.ListNasTrialDetailsRequest() ) @@ -35321,6 +35699,10 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = job_service.ListNasTrialDetailsResponse() + post_with_metadata.return_value = ( + job_service.ListNasTrialDetailsResponse(), + metadata, + ) await client.list_nas_trial_details( request, @@ -35332,6 +35714,7 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35407,10 +35790,25 @@ async def test_create_batch_prediction_job_rest_asyncio_call_success(request_typ "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, }, "input_config": { @@ -35663,10 +36061,14 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateBatchPredictionJobRequest.pb( job_service.CreateBatchPredictionJobRequest() ) @@ -35692,6 +36094,10 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = gca_batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + gca_batch_prediction_job.BatchPredictionJob(), + metadata, + ) await client.create_batch_prediction_job( request, @@ -35703,6 +36109,7 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35825,10 +36232,14 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetBatchPredictionJobRequest.pb( job_service.GetBatchPredictionJobRequest() ) @@ -35854,6 +36265,10 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + batch_prediction_job.BatchPredictionJob(), + metadata, + ) await client.get_batch_prediction_job( request, @@ -35865,6 +36280,7 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35965,10 +36381,14 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_batch_prediction_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_batch_prediction_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_batch_prediction_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListBatchPredictionJobsRequest.pb( job_service.ListBatchPredictionJobsRequest() ) @@ -35994,6 +36414,10 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = job_service.ListBatchPredictionJobsResponse() + post_with_metadata.return_value = ( + job_service.ListBatchPredictionJobsResponse(), + metadata, + ) await client.list_batch_prediction_jobs( request, @@ -36005,6 +36429,7 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36105,10 +36530,14 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteBatchPredictionJobRequest.pb( job_service.DeleteBatchPredictionJobRequest() ) @@ -36132,6 +36561,7 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_batch_prediction_job( request, @@ -36143,6 +36573,7 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36578,11 +37009,15 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_create_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateModelDeploymentMonitoringJobRequest.pb( job_service.CreateModelDeploymentMonitoringJobRequest() ) @@ -36612,6 +37047,10 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( post.return_value = ( gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) await client.create_model_deployment_monitoring_job( request, @@ -36623,6 +37062,7 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36742,11 +37182,15 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i transports.AsyncJobServiceRestInterceptor, "post_search_model_deployment_monitoring_stats_anomalies", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_search_model_deployment_monitoring_stats_anomalies_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_search_model_deployment_monitoring_stats_anomalies", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest.pb( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest() @@ -36778,6 +37222,10 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i post.return_value = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse() ) + post_with_metadata.return_value = ( + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse(), + metadata, + ) await client.search_model_deployment_monitoring_stats_anomalies( request, @@ -36789,6 +37237,7 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36923,11 +37372,15 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_get_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetModelDeploymentMonitoringJobRequest.pb( job_service.GetModelDeploymentMonitoringJobRequest() ) @@ -36957,6 +37410,10 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( post.return_value = ( model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) await client.get_model_deployment_monitoring_job( request, @@ -36968,6 +37425,7 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37075,11 +37533,15 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_list_model_deployment_monitoring_jobs", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_model_deployment_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_model_deployment_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListModelDeploymentMonitoringJobsRequest.pb( job_service.ListModelDeploymentMonitoringJobsRequest() ) @@ -37105,6 +37567,10 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = job_service.ListModelDeploymentMonitoringJobsResponse() + post_with_metadata.return_value = ( + job_service.ListModelDeploymentMonitoringJobsResponse(), + metadata, + ) await client.list_model_deployment_monitoring_jobs( request, @@ -37116,6 +37582,7 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37399,11 +37866,15 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_update_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_update_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_update_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.UpdateModelDeploymentMonitoringJobRequest.pb( job_service.UpdateModelDeploymentMonitoringJobRequest() ) @@ -37427,6 +37898,7 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_model_deployment_monitoring_job( request, @@ -37438,6 +37910,7 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37543,11 +38016,15 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_delete_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteModelDeploymentMonitoringJobRequest.pb( job_service.DeleteModelDeploymentMonitoringJobRequest() ) @@ -37571,6 +38048,7 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_deployment_monitoring_job( request, @@ -37582,6 +38060,7 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_llm_utility_service.py b/tests/unit/gapic/aiplatform_v1/test_llm_utility_service.py index 176edc7ead..c082ad9d80 100644 --- a/tests/unit/gapic/aiplatform_v1/test_llm_utility_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_llm_utility_service.py @@ -85,6 +85,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -351,6 +359,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = LlmUtilityServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = LlmUtilityServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2463,10 +2514,13 @@ def test_count_tokens_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "post_count_tokens" ) as post, mock.patch.object( + transports.LlmUtilityServiceRestInterceptor, "post_count_tokens_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "pre_count_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.CountTokensRequest.pb( prediction_service.CountTokensRequest() ) @@ -2492,6 +2546,10 @@ def test_count_tokens_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.CountTokensResponse() + post_with_metadata.return_value = ( + prediction_service.CountTokensResponse(), + metadata, + ) client.count_tokens( request, @@ -2503,6 +2561,7 @@ def test_count_tokens_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_compute_tokens_rest_bad_request( @@ -2584,10 +2643,13 @@ def test_compute_tokens_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "post_compute_tokens" ) as post, mock.patch.object( + transports.LlmUtilityServiceRestInterceptor, "post_compute_tokens_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "pre_compute_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = llm_utility_service.ComputeTokensRequest.pb( llm_utility_service.ComputeTokensRequest() ) @@ -2613,6 +2675,10 @@ def test_compute_tokens_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = llm_utility_service.ComputeTokensResponse() + post_with_metadata.return_value = ( + llm_utility_service.ComputeTokensResponse(), + metadata, + ) client.compute_tokens( request, @@ -2624,6 +2690,7 @@ def test_compute_tokens_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3409,10 +3476,14 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "post_count_tokens" ) as post, mock.patch.object( + transports.AsyncLlmUtilityServiceRestInterceptor, + "post_count_tokens_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "pre_count_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.CountTokensRequest.pb( prediction_service.CountTokensRequest() ) @@ -3438,6 +3509,10 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.CountTokensResponse() + post_with_metadata.return_value = ( + prediction_service.CountTokensResponse(), + metadata, + ) await client.count_tokens( request, @@ -3449,6 +3524,7 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3546,10 +3622,14 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "post_compute_tokens" ) as post, mock.patch.object( + transports.AsyncLlmUtilityServiceRestInterceptor, + "post_compute_tokens_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "pre_compute_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = llm_utility_service.ComputeTokensRequest.pb( llm_utility_service.ComputeTokensRequest() ) @@ -3575,6 +3655,10 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = llm_utility_service.ComputeTokensResponse() + post_with_metadata.return_value = ( + llm_utility_service.ComputeTokensResponse(), + metadata, + ) await client.compute_tokens( request, @@ -3586,6 +3670,7 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_match_service.py b/tests/unit/gapic/aiplatform_v1/test_match_service.py index 7e5c7159db..204e74495b 100644 --- a/tests/unit/gapic/aiplatform_v1/test_match_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_match_service.py @@ -76,6 +76,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -318,6 +326,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MatchServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MatchServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2095,10 +2146,13 @@ def test_find_neighbors_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MatchServiceRestInterceptor, "post_find_neighbors" ) as post, mock.patch.object( + transports.MatchServiceRestInterceptor, "post_find_neighbors_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MatchServiceRestInterceptor, "pre_find_neighbors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.FindNeighborsRequest.pb( match_service.FindNeighborsRequest() ) @@ -2124,6 +2178,10 @@ def test_find_neighbors_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.FindNeighborsResponse() + post_with_metadata.return_value = ( + match_service.FindNeighborsResponse(), + metadata, + ) client.find_neighbors( request, @@ -2135,6 +2193,7 @@ def test_find_neighbors_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_index_datapoints_rest_bad_request( @@ -2220,10 +2279,14 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MatchServiceRestInterceptor, "post_read_index_datapoints" ) as post, mock.patch.object( + transports.MatchServiceRestInterceptor, + "post_read_index_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MatchServiceRestInterceptor, "pre_read_index_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.ReadIndexDatapointsRequest.pb( match_service.ReadIndexDatapointsRequest() ) @@ -2249,6 +2312,10 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.ReadIndexDatapointsResponse() + post_with_metadata.return_value = ( + match_service.ReadIndexDatapointsResponse(), + metadata, + ) client.read_index_datapoints( request, @@ -2260,6 +2327,7 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3046,10 +3114,13 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "post_find_neighbors" ) as post, mock.patch.object( + transports.AsyncMatchServiceRestInterceptor, "post_find_neighbors_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "pre_find_neighbors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.FindNeighborsRequest.pb( match_service.FindNeighborsRequest() ) @@ -3075,6 +3146,10 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.FindNeighborsResponse() + post_with_metadata.return_value = ( + match_service.FindNeighborsResponse(), + metadata, + ) await client.find_neighbors( request, @@ -3086,6 +3161,7 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3187,10 +3263,14 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "post_read_index_datapoints" ) as post, mock.patch.object( + transports.AsyncMatchServiceRestInterceptor, + "post_read_index_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "pre_read_index_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.ReadIndexDatapointsRequest.pb( match_service.ReadIndexDatapointsRequest() ) @@ -3216,6 +3296,10 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = match_service.ReadIndexDatapointsResponse() + post_with_metadata.return_value = ( + match_service.ReadIndexDatapointsResponse(), + metadata, + ) await client.read_index_datapoints( request, @@ -3227,6 +3311,7 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_metadata_service.py b/tests/unit/gapic/aiplatform_v1/test_metadata_service.py index 982f2581f2..a660b9de03 100644 --- a/tests/unit/gapic/aiplatform_v1/test_metadata_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_metadata_service.py @@ -100,6 +100,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -357,6 +365,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MetadataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MetadataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -21569,10 +21620,14 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_create_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataStoreRequest.pb( metadata_service.CreateMetadataStoreRequest() ) @@ -21596,6 +21651,7 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_metadata_store( request, @@ -21607,6 +21663,7 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_metadata_store_rest_bad_request( @@ -21693,10 +21750,14 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_get_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataStoreRequest.pb( metadata_service.GetMetadataStoreRequest() ) @@ -21722,6 +21783,7 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_store.MetadataStore() + post_with_metadata.return_value = metadata_store.MetadataStore(), metadata client.get_metadata_store( request, @@ -21733,6 +21795,7 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_metadata_stores_rest_bad_request( @@ -21817,10 +21880,14 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_metadata_stores" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_list_metadata_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_metadata_stores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataStoresRequest.pb( metadata_service.ListMetadataStoresRequest() ) @@ -21846,6 +21913,10 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataStoresResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataStoresResponse(), + metadata, + ) client.list_metadata_stores( request, @@ -21857,6 +21928,7 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_metadata_store_rest_bad_request( @@ -21937,10 +22009,14 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_delete_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteMetadataStoreRequest.pb( metadata_service.DeleteMetadataStoreRequest() ) @@ -21964,6 +22040,7 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_metadata_store( request, @@ -21975,6 +22052,7 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_artifact_rest_bad_request( @@ -22158,10 +22236,13 @@ def test_create_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateArtifactRequest.pb( metadata_service.CreateArtifactRequest() ) @@ -22185,6 +22266,7 @@ def test_create_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata client.create_artifact( request, @@ -22196,6 +22278,7 @@ def test_create_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_artifact_rest_bad_request( @@ -22298,10 +22381,13 @@ def test_get_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetArtifactRequest.pb( metadata_service.GetArtifactRequest() ) @@ -22325,6 +22411,7 @@ def test_get_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = artifact.Artifact() + post_with_metadata.return_value = artifact.Artifact(), metadata client.get_artifact( request, @@ -22336,6 +22423,7 @@ def test_get_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_artifacts_rest_bad_request( @@ -22424,10 +22512,13 @@ def test_list_artifacts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_artifacts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_artifacts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListArtifactsRequest.pb( metadata_service.ListArtifactsRequest() ) @@ -22453,6 +22544,10 @@ def test_list_artifacts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListArtifactsResponse() + post_with_metadata.return_value = ( + metadata_service.ListArtifactsResponse(), + metadata, + ) client.list_artifacts( request, @@ -22464,6 +22559,7 @@ def test_list_artifacts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_artifact_rest_bad_request( @@ -22651,10 +22747,13 @@ def test_update_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateArtifactRequest.pb( metadata_service.UpdateArtifactRequest() ) @@ -22678,6 +22777,7 @@ def test_update_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata client.update_artifact( request, @@ -22689,6 +22789,7 @@ def test_update_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_artifact_rest_bad_request( @@ -22773,10 +22874,13 @@ def test_delete_artifact_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteArtifactRequest.pb( metadata_service.DeleteArtifactRequest() ) @@ -22800,6 +22904,7 @@ def test_delete_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_artifact( request, @@ -22811,6 +22916,7 @@ def test_delete_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_artifacts_rest_bad_request( @@ -22895,10 +23001,13 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_artifacts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_artifacts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeArtifactsRequest.pb( metadata_service.PurgeArtifactsRequest() ) @@ -22922,6 +23031,7 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_artifacts( request, @@ -22933,6 +23043,7 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_context_rest_bad_request( @@ -23113,10 +23224,13 @@ def test_create_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateContextRequest.pb( metadata_service.CreateContextRequest() ) @@ -23140,6 +23254,7 @@ def test_create_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata client.create_context( request, @@ -23151,6 +23266,7 @@ def test_create_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_context_rest_bad_request(request_type=metadata_service.GetContextRequest): @@ -23249,10 +23365,13 @@ def test_get_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetContextRequest.pb( metadata_service.GetContextRequest() ) @@ -23276,6 +23395,7 @@ def test_get_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = context.Context() + post_with_metadata.return_value = context.Context(), metadata client.get_context( request, @@ -23287,6 +23407,7 @@ def test_get_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_contexts_rest_bad_request( @@ -23375,10 +23496,13 @@ def test_list_contexts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_contexts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_contexts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListContextsRequest.pb( metadata_service.ListContextsRequest() ) @@ -23404,6 +23528,10 @@ def test_list_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListContextsResponse() + post_with_metadata.return_value = ( + metadata_service.ListContextsResponse(), + metadata, + ) client.list_contexts( request, @@ -23415,6 +23543,7 @@ def test_list_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_context_rest_bad_request( @@ -23599,10 +23728,13 @@ def test_update_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateContextRequest.pb( metadata_service.UpdateContextRequest() ) @@ -23626,6 +23758,7 @@ def test_update_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata client.update_context( request, @@ -23637,6 +23770,7 @@ def test_update_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_context_rest_bad_request( @@ -23721,10 +23855,13 @@ def test_delete_context_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteContextRequest.pb( metadata_service.DeleteContextRequest() ) @@ -23748,6 +23885,7 @@ def test_delete_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_context( request, @@ -23759,6 +23897,7 @@ def test_delete_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_contexts_rest_bad_request( @@ -23843,10 +23982,13 @@ def test_purge_contexts_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_contexts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_contexts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeContextsRequest.pb( metadata_service.PurgeContextsRequest() ) @@ -23870,6 +24012,7 @@ def test_purge_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_contexts( request, @@ -23881,6 +24024,7 @@ def test_purge_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_context_artifacts_and_executions_rest_bad_request( @@ -23971,11 +24115,15 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor transports.MetadataServiceRestInterceptor, "post_add_context_artifacts_and_executions", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_context_artifacts_and_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_context_artifacts_and_executions", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextArtifactsAndExecutionsRequest.pb( metadata_service.AddContextArtifactsAndExecutionsRequest() ) @@ -24003,6 +24151,10 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextArtifactsAndExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextArtifactsAndExecutionsResponse(), + metadata, + ) client.add_context_artifacts_and_executions( request, @@ -24014,6 +24166,7 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_context_children_rest_bad_request( @@ -24099,10 +24252,14 @@ def test_add_context_children_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_add_context_children" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextChildrenRequest.pb( metadata_service.AddContextChildrenRequest() ) @@ -24128,6 +24285,10 @@ def test_add_context_children_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextChildrenResponse(), + metadata, + ) client.add_context_children( request, @@ -24139,6 +24300,7 @@ def test_add_context_children_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_remove_context_children_rest_bad_request( @@ -24224,10 +24386,14 @@ def test_remove_context_children_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_remove_context_children" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_remove_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_remove_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.RemoveContextChildrenRequest.pb( metadata_service.RemoveContextChildrenRequest() ) @@ -24253,6 +24419,10 @@ def test_remove_context_children_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.RemoveContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.RemoveContextChildrenResponse(), + metadata, + ) client.remove_context_children( request, @@ -24264,6 +24434,7 @@ def test_remove_context_children_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_context_lineage_subgraph_rest_bad_request( @@ -24349,10 +24520,14 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_query_context_lineage_subgraph" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_context_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_context_lineage_subgraph" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryContextLineageSubgraphRequest.pb( metadata_service.QueryContextLineageSubgraphRequest() ) @@ -24378,6 +24553,7 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_context_lineage_subgraph( request, @@ -24389,6 +24565,7 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_execution_rest_bad_request( @@ -24569,10 +24746,13 @@ def test_create_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateExecutionRequest.pb( metadata_service.CreateExecutionRequest() ) @@ -24596,6 +24776,7 @@ def test_create_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata client.create_execution( request, @@ -24607,6 +24788,7 @@ def test_create_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_execution_rest_bad_request( @@ -24707,10 +24889,13 @@ def test_get_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetExecutionRequest.pb( metadata_service.GetExecutionRequest() ) @@ -24734,6 +24919,7 @@ def test_get_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = execution.Execution() + post_with_metadata.return_value = execution.Execution(), metadata client.get_execution( request, @@ -24745,6 +24931,7 @@ def test_get_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_executions_rest_bad_request( @@ -24833,10 +25020,13 @@ def test_list_executions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_executions" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_executions_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListExecutionsRequest.pb( metadata_service.ListExecutionsRequest() ) @@ -24862,6 +25052,10 @@ def test_list_executions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.ListExecutionsResponse(), + metadata, + ) client.list_executions( request, @@ -24873,6 +25067,7 @@ def test_list_executions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_execution_rest_bad_request( @@ -25057,10 +25252,13 @@ def test_update_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateExecutionRequest.pb( metadata_service.UpdateExecutionRequest() ) @@ -25084,6 +25282,7 @@ def test_update_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata client.update_execution( request, @@ -25095,6 +25294,7 @@ def test_update_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_execution_rest_bad_request( @@ -25179,10 +25379,13 @@ def test_delete_execution_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteExecutionRequest.pb( metadata_service.DeleteExecutionRequest() ) @@ -25206,6 +25409,7 @@ def test_delete_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_execution( request, @@ -25217,6 +25421,7 @@ def test_delete_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_executions_rest_bad_request( @@ -25301,10 +25506,13 @@ def test_purge_executions_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_executions" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_executions_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeExecutionsRequest.pb( metadata_service.PurgeExecutionsRequest() ) @@ -25328,6 +25536,7 @@ def test_purge_executions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_executions( request, @@ -25339,6 +25548,7 @@ def test_purge_executions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_execution_events_rest_bad_request( @@ -25424,10 +25634,14 @@ def test_add_execution_events_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_add_execution_events" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_execution_events_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_execution_events" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddExecutionEventsRequest.pb( metadata_service.AddExecutionEventsRequest() ) @@ -25453,6 +25667,10 @@ def test_add_execution_events_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddExecutionEventsResponse() + post_with_metadata.return_value = ( + metadata_service.AddExecutionEventsResponse(), + metadata, + ) client.add_execution_events( request, @@ -25464,6 +25682,7 @@ def test_add_execution_events_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_execution_inputs_and_outputs_rest_bad_request( @@ -25550,11 +25769,15 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): transports.MetadataServiceRestInterceptor, "post_query_execution_inputs_and_outputs", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_execution_inputs_and_outputs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_execution_inputs_and_outputs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryExecutionInputsAndOutputsRequest.pb( metadata_service.QueryExecutionInputsAndOutputsRequest() ) @@ -25580,6 +25803,7 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_execution_inputs_and_outputs( request, @@ -25591,6 +25815,7 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_metadata_schema_rest_bad_request( @@ -25767,10 +25992,14 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_metadata_schema" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_create_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataSchemaRequest.pb( metadata_service.CreateMetadataSchemaRequest() ) @@ -25796,6 +26025,7 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_metadata_schema.MetadataSchema() + post_with_metadata.return_value = gca_metadata_schema.MetadataSchema(), metadata client.create_metadata_schema( request, @@ -25807,6 +26037,7 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_metadata_schema_rest_bad_request( @@ -25906,10 +26137,14 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_metadata_schema" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_get_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataSchemaRequest.pb( metadata_service.GetMetadataSchemaRequest() ) @@ -25935,6 +26170,7 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_schema.MetadataSchema() + post_with_metadata.return_value = metadata_schema.MetadataSchema(), metadata client.get_metadata_schema( request, @@ -25946,6 +26182,7 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_metadata_schemas_rest_bad_request( @@ -26034,10 +26271,14 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_metadata_schemas" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_list_metadata_schemas_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_metadata_schemas" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataSchemasRequest.pb( metadata_service.ListMetadataSchemasRequest() ) @@ -26063,6 +26304,10 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataSchemasResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataSchemasResponse(), + metadata, + ) client.list_metadata_schemas( request, @@ -26074,6 +26319,7 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_artifact_lineage_subgraph_rest_bad_request( @@ -26160,10 +26406,14 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): transports.MetadataServiceRestInterceptor, "post_query_artifact_lineage_subgraph", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_artifact_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_artifact_lineage_subgraph" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryArtifactLineageSubgraphRequest.pb( metadata_service.QueryArtifactLineageSubgraphRequest() ) @@ -26189,6 +26439,7 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_artifact_lineage_subgraph( request, @@ -26200,6 +26451,7 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -27702,10 +27954,14 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataStoreRequest.pb( metadata_service.CreateMetadataStoreRequest() ) @@ -27729,6 +27985,7 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_metadata_store( request, @@ -27740,6 +27997,7 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27842,10 +28100,14 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataStoreRequest.pb( metadata_service.GetMetadataStoreRequest() ) @@ -27871,6 +28133,7 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_store.MetadataStore() + post_with_metadata.return_value = metadata_store.MetadataStore(), metadata await client.get_metadata_store( request, @@ -27882,6 +28145,7 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27982,10 +28246,14 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_metadata_stores" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_metadata_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_metadata_stores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataStoresRequest.pb( metadata_service.ListMetadataStoresRequest() ) @@ -28011,6 +28279,10 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataStoresResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataStoresResponse(), + metadata, + ) await client.list_metadata_stores( request, @@ -28022,6 +28294,7 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28118,10 +28391,14 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteMetadataStoreRequest.pb( metadata_service.DeleteMetadataStoreRequest() ) @@ -28145,6 +28422,7 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_metadata_store( request, @@ -28156,6 +28434,7 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28355,10 +28634,14 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateArtifactRequest.pb( metadata_service.CreateArtifactRequest() ) @@ -28382,6 +28665,7 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata await client.create_artifact( request, @@ -28393,6 +28677,7 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28511,10 +28796,14 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetArtifactRequest.pb( metadata_service.GetArtifactRequest() ) @@ -28538,6 +28827,7 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = artifact.Artifact() + post_with_metadata.return_value = artifact.Artifact(), metadata await client.get_artifact( request, @@ -28549,6 +28839,7 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28653,10 +28944,14 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_artifacts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_artifacts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListArtifactsRequest.pb( metadata_service.ListArtifactsRequest() ) @@ -28682,6 +28977,10 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListArtifactsResponse() + post_with_metadata.return_value = ( + metadata_service.ListArtifactsResponse(), + metadata, + ) await client.list_artifacts( request, @@ -28693,6 +28992,7 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28896,10 +29196,14 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateArtifactRequest.pb( metadata_service.UpdateArtifactRequest() ) @@ -28923,6 +29227,7 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata await client.update_artifact( request, @@ -28934,6 +29239,7 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29034,10 +29340,14 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteArtifactRequest.pb( metadata_service.DeleteArtifactRequest() ) @@ -29061,6 +29371,7 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_artifact( request, @@ -29072,6 +29383,7 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29172,10 +29484,14 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_artifacts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_artifacts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeArtifactsRequest.pb( metadata_service.PurgeArtifactsRequest() ) @@ -29199,6 +29515,7 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_artifacts( request, @@ -29210,6 +29527,7 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29406,10 +29724,14 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateContextRequest.pb( metadata_service.CreateContextRequest() ) @@ -29433,6 +29755,7 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata await client.create_context( request, @@ -29444,6 +29767,7 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29560,10 +29884,13 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, "post_get_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetContextRequest.pb( metadata_service.GetContextRequest() ) @@ -29587,6 +29914,7 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = context.Context() + post_with_metadata.return_value = context.Context(), metadata await client.get_context( request, @@ -29598,6 +29926,7 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29702,10 +30031,14 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_contexts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListContextsRequest.pb( metadata_service.ListContextsRequest() ) @@ -29731,6 +30064,10 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListContextsResponse() + post_with_metadata.return_value = ( + metadata_service.ListContextsResponse(), + metadata, + ) await client.list_contexts( request, @@ -29742,6 +30079,7 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29942,10 +30280,14 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateContextRequest.pb( metadata_service.UpdateContextRequest() ) @@ -29969,6 +30311,7 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata await client.update_context( request, @@ -29980,6 +30323,7 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30080,10 +30424,14 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteContextRequest.pb( metadata_service.DeleteContextRequest() ) @@ -30107,6 +30455,7 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_context( request, @@ -30118,6 +30467,7 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30218,10 +30568,14 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_contexts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeContextsRequest.pb( metadata_service.PurgeContextsRequest() ) @@ -30245,6 +30599,7 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_contexts( request, @@ -30256,6 +30611,7 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30366,11 +30722,15 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_add_context_artifacts_and_executions", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_context_artifacts_and_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_context_artifacts_and_executions", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextArtifactsAndExecutionsRequest.pb( metadata_service.AddContextArtifactsAndExecutionsRequest() ) @@ -30398,6 +30758,10 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextArtifactsAndExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextArtifactsAndExecutionsResponse(), + metadata, + ) await client.add_context_artifacts_and_executions( request, @@ -30409,6 +30773,7 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30510,10 +30875,14 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_add_context_children" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextChildrenRequest.pb( metadata_service.AddContextChildrenRequest() ) @@ -30539,6 +30908,10 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextChildrenResponse(), + metadata, + ) await client.add_context_children( request, @@ -30550,6 +30923,7 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30651,10 +31025,14 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_remove_context_children" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_remove_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_remove_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.RemoveContextChildrenRequest.pb( metadata_service.RemoveContextChildrenRequest() ) @@ -30680,6 +31058,10 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = metadata_service.RemoveContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.RemoveContextChildrenResponse(), + metadata, + ) await client.remove_context_children( request, @@ -30691,6 +31073,7 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30795,11 +31178,15 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_context_lineage_subgraph", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_context_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_context_lineage_subgraph", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryContextLineageSubgraphRequest.pb( metadata_service.QueryContextLineageSubgraphRequest() ) @@ -30825,6 +31212,7 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_context_lineage_subgraph( request, @@ -30836,6 +31224,7 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31032,10 +31421,14 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateExecutionRequest.pb( metadata_service.CreateExecutionRequest() ) @@ -31059,6 +31452,7 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata await client.create_execution( request, @@ -31070,6 +31464,7 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31186,10 +31581,14 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetExecutionRequest.pb( metadata_service.GetExecutionRequest() ) @@ -31213,6 +31612,7 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = execution.Execution() + post_with_metadata.return_value = execution.Execution(), metadata await client.get_execution( request, @@ -31224,6 +31624,7 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31328,10 +31729,14 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_executions" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListExecutionsRequest.pb( metadata_service.ListExecutionsRequest() ) @@ -31357,6 +31762,10 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.ListExecutionsResponse(), + metadata, + ) await client.list_executions( request, @@ -31368,6 +31777,7 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31568,10 +31978,14 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateExecutionRequest.pb( metadata_service.UpdateExecutionRequest() ) @@ -31595,6 +32009,7 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata await client.update_execution( request, @@ -31606,6 +32021,7 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31706,10 +32122,14 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteExecutionRequest.pb( metadata_service.DeleteExecutionRequest() ) @@ -31733,6 +32153,7 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_execution( request, @@ -31744,6 +32165,7 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31844,10 +32266,14 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_executions" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeExecutionsRequest.pb( metadata_service.PurgeExecutionsRequest() ) @@ -31871,6 +32297,7 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_executions( request, @@ -31882,6 +32309,7 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31983,10 +32411,14 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_add_execution_events" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_execution_events_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_execution_events" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddExecutionEventsRequest.pb( metadata_service.AddExecutionEventsRequest() ) @@ -32012,6 +32444,10 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddExecutionEventsResponse() + post_with_metadata.return_value = ( + metadata_service.AddExecutionEventsResponse(), + metadata, + ) await client.add_execution_events( request, @@ -32023,6 +32459,7 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32129,11 +32566,15 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_execution_inputs_and_outputs", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_execution_inputs_and_outputs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_execution_inputs_and_outputs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryExecutionInputsAndOutputsRequest.pb( metadata_service.QueryExecutionInputsAndOutputsRequest() ) @@ -32159,6 +32600,7 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_execution_inputs_and_outputs( request, @@ -32170,6 +32612,7 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32362,10 +32805,14 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_metadata_schema" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataSchemaRequest.pb( metadata_service.CreateMetadataSchemaRequest() ) @@ -32391,6 +32838,7 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_metadata_schema.MetadataSchema() + post_with_metadata.return_value = gca_metadata_schema.MetadataSchema(), metadata await client.create_metadata_schema( request, @@ -32402,6 +32850,7 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32517,10 +32966,14 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_metadata_schema" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataSchemaRequest.pb( metadata_service.GetMetadataSchemaRequest() ) @@ -32546,6 +32999,7 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_schema.MetadataSchema() + post_with_metadata.return_value = metadata_schema.MetadataSchema(), metadata await client.get_metadata_schema( request, @@ -32557,6 +33011,7 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32661,10 +33116,14 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_metadata_schemas" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_metadata_schemas_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_metadata_schemas" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataSchemasRequest.pb( metadata_service.ListMetadataSchemasRequest() ) @@ -32690,6 +33149,10 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataSchemasResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataSchemasResponse(), + metadata, + ) await client.list_metadata_schemas( request, @@ -32701,6 +33164,7 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32805,11 +33269,15 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_artifact_lineage_subgraph", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_artifact_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_artifact_lineage_subgraph", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryArtifactLineageSubgraphRequest.pb( metadata_service.QueryArtifactLineageSubgraphRequest() ) @@ -32835,6 +33303,7 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_artifact_lineage_subgraph( request, @@ -32846,6 +33315,7 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_migration_service.py b/tests/unit/gapic/aiplatform_v1/test_migration_service.py index 56df198495..0b26e94545 100644 --- a/tests/unit/gapic/aiplatform_v1/test_migration_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_migration_service.py @@ -83,6 +83,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -341,6 +349,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MigrationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MigrationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2832,10 +2883,14 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MigrationServiceRestInterceptor, "post_search_migratable_resources" ) as post, mock.patch.object( + transports.MigrationServiceRestInterceptor, + "post_search_migratable_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MigrationServiceRestInterceptor, "pre_search_migratable_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.SearchMigratableResourcesRequest.pb( migration_service.SearchMigratableResourcesRequest() ) @@ -2861,6 +2916,10 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = migration_service.SearchMigratableResourcesResponse() + post_with_metadata.return_value = ( + migration_service.SearchMigratableResourcesResponse(), + metadata, + ) client.search_migratable_resources( request, @@ -2872,6 +2931,7 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_migrate_resources_rest_bad_request( @@ -2952,10 +3012,14 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MigrationServiceRestInterceptor, "post_batch_migrate_resources" ) as post, mock.patch.object( + transports.MigrationServiceRestInterceptor, + "post_batch_migrate_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MigrationServiceRestInterceptor, "pre_batch_migrate_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.BatchMigrateResourcesRequest.pb( migration_service.BatchMigrateResourcesRequest() ) @@ -2979,6 +3043,7 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_migrate_resources( request, @@ -2990,6 +3055,7 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3797,11 +3863,15 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc transports.AsyncMigrationServiceRestInterceptor, "post_search_migratable_resources", ) as post, mock.patch.object( + transports.AsyncMigrationServiceRestInterceptor, + "post_search_migratable_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "pre_search_migratable_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.SearchMigratableResourcesRequest.pb( migration_service.SearchMigratableResourcesRequest() ) @@ -3827,6 +3897,10 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = migration_service.SearchMigratableResourcesResponse() + post_with_metadata.return_value = ( + migration_service.SearchMigratableResourcesResponse(), + metadata, + ) await client.search_migratable_resources( request, @@ -3838,6 +3912,7 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3934,10 +4009,14 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "post_batch_migrate_resources" ) as post, mock.patch.object( + transports.AsyncMigrationServiceRestInterceptor, + "post_batch_migrate_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "pre_batch_migrate_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.BatchMigrateResourcesRequest.pb( migration_service.BatchMigrateResourcesRequest() ) @@ -3961,6 +4040,7 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_migrate_resources( request, @@ -3972,6 +4052,7 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -5317,22 +5398,19 @@ def test_parse_annotated_dataset_path(): def test_dataset_path(): project = "cuttlefish" - location = "mussel" - dataset = "winkle" - expected = "projects/{project}/locations/{location}/datasets/{dataset}".format( + dataset = "mussel" + expected = "projects/{project}/datasets/{dataset}".format( project=project, - location=location, dataset=dataset, ) - actual = MigrationServiceClient.dataset_path(project, location, dataset) + actual = MigrationServiceClient.dataset_path(project, dataset) assert expected == actual def test_parse_dataset_path(): expected = { - "project": "nautilus", - "location": "scallop", - "dataset": "abalone", + "project": "winkle", + "dataset": "nautilus", } path = MigrationServiceClient.dataset_path(**expected) @@ -5342,19 +5420,22 @@ def test_parse_dataset_path(): def test_dataset_path(): - project = "squid" - dataset = "clam" - expected = "projects/{project}/datasets/{dataset}".format( + project = "scallop" + location = "abalone" + dataset = "squid" + expected = "projects/{project}/locations/{location}/datasets/{dataset}".format( project=project, + location=location, dataset=dataset, ) - actual = MigrationServiceClient.dataset_path(project, dataset) + actual = MigrationServiceClient.dataset_path(project, location, dataset) assert expected == actual def test_parse_dataset_path(): expected = { - "project": "whelk", + "project": "clam", + "location": "whelk", "dataset": "octopus", } path = MigrationServiceClient.dataset_path(**expected) diff --git a/tests/unit/gapic/aiplatform_v1/test_model_garden_service.py b/tests/unit/gapic/aiplatform_v1/test_model_garden_service.py index 3bf3db409b..dd2f57bb64 100644 --- a/tests/unit/gapic/aiplatform_v1/test_model_garden_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_model_garden_service.py @@ -81,6 +81,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -351,6 +359,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ModelGardenServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ModelGardenServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1994,10 +2045,14 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "post_get_publisher_model" ) as post, mock.patch.object( + transports.ModelGardenServiceRestInterceptor, + "post_get_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "pre_get_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.GetPublisherModelRequest.pb( model_garden_service.GetPublisherModelRequest() ) @@ -2023,6 +2078,7 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = publisher_model.PublisherModel() + post_with_metadata.return_value = publisher_model.PublisherModel(), metadata client.get_publisher_model( request, @@ -2034,6 +2090,7 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -2819,10 +2876,14 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "post_get_publisher_model" ) as post, mock.patch.object( + transports.AsyncModelGardenServiceRestInterceptor, + "post_get_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "pre_get_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.GetPublisherModelRequest.pb( model_garden_service.GetPublisherModelRequest() ) @@ -2848,6 +2909,7 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = publisher_model.PublisherModel() + post_with_metadata.return_value = publisher_model.PublisherModel(), metadata await client.get_publisher_model( request, @@ -2859,6 +2921,7 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_model_service.py b/tests/unit/gapic/aiplatform_v1/test_model_service.py index fbf212a5f5..420f8f1837 100644 --- a/tests/unit/gapic/aiplatform_v1/test_model_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_model_service.py @@ -99,6 +99,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -341,6 +349,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ModelServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ModelServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1445,6 +1496,7 @@ def test_get_model(request_type, transport: str = "grpc"): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -1475,6 +1527,7 @@ def test_get_model(request_type, transport: str = "grpc"): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -1621,6 +1674,7 @@ async def test_get_model_async( display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -1656,6 +1710,7 @@ async def test_get_model_async( assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -2894,11 +2949,11 @@ async def test_list_model_versions_async_pages(): @pytest.mark.parametrize( "request_type", [ - model_service.UpdateModelRequest, + model_service.ListModelVersionCheckpointsRequest, dict, ], ) -def test_update_model(request_type, transport: str = "grpc"): +def test_list_model_version_checkpoints(request_type, transport: str = "grpc"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -2909,65 +2964,27 @@ def test_update_model(request_type, transport: str = "grpc"): request = request_type() # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model( - name="name_value", - version_id="version_id_value", - version_aliases=["version_aliases_value"], - display_name="display_name_value", - description="description_value", - version_description="version_description_value", - metadata_schema_uri="metadata_schema_uri_value", - training_pipeline="training_pipeline_value", - pipeline_job="pipeline_job_value", - artifact_uri="artifact_uri_value", - supported_deployment_resources_types=[ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ], - supported_input_storage_formats=["supported_input_storage_formats_value"], - supported_output_storage_formats=["supported_output_storage_formats_value"], - etag="etag_value", - metadata_artifact="metadata_artifact_value", - satisfies_pzs=True, - satisfies_pzi=True, + call.return_value = model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", ) - response = client.update_model(request) + response = client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) == 1 _, args, _ = call.mock_calls[0] - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() assert args[0] == request # Establish that the response is the type that we expect. - assert isinstance(response, gca_model.Model) - assert response.name == "name_value" - assert response.version_id == "version_id_value" - assert response.version_aliases == ["version_aliases_value"] - assert response.display_name == "display_name_value" - assert response.description == "description_value" - assert response.version_description == "version_description_value" - assert response.metadata_schema_uri == "metadata_schema_uri_value" - assert response.training_pipeline == "training_pipeline_value" - assert response.pipeline_job == "pipeline_job_value" - assert response.artifact_uri == "artifact_uri_value" - assert response.supported_deployment_resources_types == [ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ] - assert response.supported_input_storage_formats == [ - "supported_input_storage_formats_value" - ] - assert response.supported_output_storage_formats == [ - "supported_output_storage_formats_value" - ] - assert response.etag == "etag_value" - assert response.metadata_artifact == "metadata_artifact_value" - assert response.satisfies_pzs is True - assert response.satisfies_pzi is True + assert isinstance(response, pagers.ListModelVersionCheckpointsPager) + assert response.next_page_token == "next_page_token_value" -def test_update_model_non_empty_request_with_auto_populated_field(): +def test_list_model_version_checkpoints_non_empty_request_with_auto_populated_field(): # This test is a coverage failsafe to make sure that UUID4 fields are # automatically populated, according to AIP-4235, with non-empty requests. client = ModelServiceClient( @@ -2978,20 +2995,28 @@ def test_update_model_non_empty_request_with_auto_populated_field(): # Populate all string fields in the request which are not UUID4 # since we want to check that UUID4 are populated automatically # if they meet the requirements of AIP 4235. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest( + name="name_value", + page_token="page_token_value", + ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: call.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client.update_model(request=request) + client.list_model_version_checkpoints(request=request) call.assert_called() _, args, _ = call.mock_calls[0] - assert args[0] == model_service.UpdateModelRequest() + assert args[0] == model_service.ListModelVersionCheckpointsRequest( + name="name_value", + page_token="page_token_value", + ) -def test_update_model_use_cached_wrapped_rpc(): +def test_list_model_version_checkpoints_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -3005,21 +3030,26 @@ def test_update_model_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.update_model in client._transport._wrapped_methods + assert ( + client._transport.list_model_version_checkpoints + in client._transport._wrapped_methods + ) # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.update_model] = mock_rpc + client._transport._wrapped_methods[ + client._transport.list_model_version_checkpoints + ] = mock_rpc request = {} - client.update_model(request) + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.update_model(request) + client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 @@ -3027,7 +3057,7 @@ def test_update_model_use_cached_wrapped_rpc(): @pytest.mark.asyncio -async def test_update_model_async_use_cached_wrapped_rpc( +async def test_list_model_version_checkpoints_async_use_cached_wrapped_rpc( transport: str = "grpc_asyncio", ): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, @@ -3044,7 +3074,7 @@ async def test_update_model_async_use_cached_wrapped_rpc( # Ensure method has been cached assert ( - client._client._transport.update_model + client._client._transport.list_model_version_checkpoints in client._client._transport._wrapped_methods ) @@ -3052,16 +3082,16 @@ async def test_update_model_async_use_cached_wrapped_rpc( mock_rpc = mock.AsyncMock() mock_rpc.return_value = mock.Mock() client._client._transport._wrapped_methods[ - client._client._transport.update_model + client._client._transport.list_model_version_checkpoints ] = mock_rpc request = {} - await client.update_model(request) + await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - await client.update_model(request) + await client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 @@ -3069,8 +3099,9 @@ async def test_update_model_async_use_cached_wrapped_rpc( @pytest.mark.asyncio -async def test_update_model_async( - transport: str = "grpc_asyncio", request_type=model_service.UpdateModelRequest +async def test_list_model_version_checkpoints_async( + transport: str = "grpc_asyncio", + request_type=model_service.ListModelVersionCheckpointsRequest, ): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), @@ -3082,90 +3113,50 @@ async def test_update_model_async( request = request_type() # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( - gca_model.Model( - name="name_value", - version_id="version_id_value", - version_aliases=["version_aliases_value"], - display_name="display_name_value", - description="description_value", - version_description="version_description_value", - metadata_schema_uri="metadata_schema_uri_value", - training_pipeline="training_pipeline_value", - pipeline_job="pipeline_job_value", - artifact_uri="artifact_uri_value", - supported_deployment_resources_types=[ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ], - supported_input_storage_formats=[ - "supported_input_storage_formats_value" - ], - supported_output_storage_formats=[ - "supported_output_storage_formats_value" - ], - etag="etag_value", - metadata_artifact="metadata_artifact_value", - satisfies_pzs=True, - satisfies_pzi=True, + model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", ) ) - response = await client.update_model(request) + response = await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) _, args, _ = call.mock_calls[0] - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() assert args[0] == request # Establish that the response is the type that we expect. - assert isinstance(response, gca_model.Model) - assert response.name == "name_value" - assert response.version_id == "version_id_value" - assert response.version_aliases == ["version_aliases_value"] - assert response.display_name == "display_name_value" - assert response.description == "description_value" - assert response.version_description == "version_description_value" - assert response.metadata_schema_uri == "metadata_schema_uri_value" - assert response.training_pipeline == "training_pipeline_value" - assert response.pipeline_job == "pipeline_job_value" - assert response.artifact_uri == "artifact_uri_value" - assert response.supported_deployment_resources_types == [ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ] - assert response.supported_input_storage_formats == [ - "supported_input_storage_formats_value" - ] - assert response.supported_output_storage_formats == [ - "supported_output_storage_formats_value" - ] - assert response.etag == "etag_value" - assert response.metadata_artifact == "metadata_artifact_value" - assert response.satisfies_pzs is True - assert response.satisfies_pzi is True + assert isinstance(response, pagers.ListModelVersionCheckpointsAsyncPager) + assert response.next_page_token == "next_page_token_value" @pytest.mark.asyncio -async def test_update_model_async_from_dict(): - await test_update_model_async(request_type=dict) +async def test_list_model_version_checkpoints_async_from_dict(): + await test_list_model_version_checkpoints_async(request_type=dict) -def test_update_model_field_headers(): +def test_list_model_version_checkpoints_field_headers(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) # Any value that is part of the HTTP/1.1 URI should be sent as # a field header. Set these to a non-empty value. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() - request.model.name = "name_value" + request.name = "name_value" # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: - call.return_value = gca_model.Model() - client.update_model(request) + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = model_service.ListModelVersionCheckpointsResponse() + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) == 1 @@ -3176,26 +3167,30 @@ def test_update_model_field_headers(): _, _, kw = call.mock_calls[0] assert ( "x-goog-request-params", - "model.name=name_value", + "name=name_value", ) in kw["metadata"] @pytest.mark.asyncio -async def test_update_model_field_headers_async(): +async def test_list_model_version_checkpoints_field_headers_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) # Any value that is part of the HTTP/1.1 URI should be sent as # a field header. Set these to a non-empty value. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() - request.model.name = "name_value" + request.name = "name_value" # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: - call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) - await client.update_model(request) + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse() + ) + await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) @@ -3206,39 +3201,37 @@ async def test_update_model_field_headers_async(): _, _, kw = call.mock_calls[0] assert ( "x-goog-request-params", - "model.name=name_value", + "name=name_value", ) in kw["metadata"] -def test_update_model_flattened(): +def test_list_model_version_checkpoints_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model() + call.return_value = model_service.ListModelVersionCheckpointsResponse() # Call the method with a truthy value for each flattened field, # using the keyword arguments to the method. - client.update_model( - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + client.list_model_version_checkpoints( + name="name_value", ) # Establish that the underlying call was made with the expected # request object values. assert len(call.mock_calls) == 1 _, args, _ = call.mock_calls[0] - arg = args[0].model - mock_val = gca_model.Model(name="name_value") - assert arg == mock_val - arg = args[0].update_mask - mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + arg = args[0].name + mock_val = "name_value" assert arg == mock_val -def test_update_model_flattened_error(): +def test_list_model_version_checkpoints_flattened_error(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) @@ -3246,46 +3239,45 @@ def test_update_model_flattened_error(): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.update_model( - model_service.UpdateModelRequest(), - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), + name="name_value", ) @pytest.mark.asyncio -async def test_update_model_flattened_async(): +async def test_list_model_version_checkpoints_flattened_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model() + call.return_value = model_service.ListModelVersionCheckpointsResponse() - call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse() + ) # Call the method with a truthy value for each flattened field, # using the keyword arguments to the method. - response = await client.update_model( - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + response = await client.list_model_version_checkpoints( + name="name_value", ) # Establish that the underlying call was made with the expected # request object values. assert len(call.mock_calls) _, args, _ = call.mock_calls[0] - arg = args[0].model - mock_val = gca_model.Model(name="name_value") - assert arg == mock_val - arg = args[0].update_mask - mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + arg = args[0].name + mock_val = "name_value" assert arg == mock_val @pytest.mark.asyncio -async def test_update_model_flattened_error_async(): +async def test_list_model_version_checkpoints_flattened_error_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) @@ -3293,50 +3285,668 @@ async def test_update_model_flattened_error_async(): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - await client.update_model( - model_service.UpdateModelRequest(), - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + await client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), + name="name_value", ) -@pytest.mark.parametrize( - "request_type", - [ - model_service.UpdateExplanationDatasetRequest, - dict, - ], -) -def test_update_explanation_dataset(request_type, transport: str = "grpc"): +def test_list_model_version_checkpoints_pager(transport_name: str = "grpc"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), - transport=transport, + transport=transport_name, ) - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = request_type() - # Mock the actual call within the gRPC stub, and fake the request. with mock.patch.object( - type(client.transport.update_explanation_dataset), "__call__" + type(client.transport.list_model_version_checkpoints), "__call__" ) as call: - # Designate an appropriate return value for the call. - call.return_value = operations_pb2.Operation(name="operations/spam") - response = client.update_explanation_dataset(request) - - # Establish that the underlying gRPC stub method was called. - assert len(call.mock_calls) == 1 - _, args, _ = call.mock_calls[0] - request = model_service.UpdateExplanationDatasetRequest() - assert args[0] == request - - # Establish that the response is the type that we expect. - assert isinstance(response, future.Future) - - -def test_update_explanation_dataset_non_empty_request_with_auto_populated_field(): - # This test is a coverage failsafe to make sure that UUID4 fields are + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + + expected_metadata = () + retry = retries.Retry() + timeout = 5 + expected_metadata = tuple(expected_metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", ""),)), + ) + pager = client.list_model_version_checkpoints( + request={}, retry=retry, timeout=timeout + ) + + assert pager._metadata == expected_metadata + assert pager._retry == retry + assert pager._timeout == timeout + + results = list(pager) + assert len(results) == 6 + assert all(isinstance(i, model_service.ModelVersionCheckpoint) for i in results) + + +def test_list_model_version_checkpoints_pages(transport_name: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport_name, + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + pages = list(client.list_model_version_checkpoints(request={}).pages) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_async_pager(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), + "__call__", + new_callable=mock.AsyncMock, + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + async_pager = await client.list_model_version_checkpoints( + request={}, + ) + assert async_pager.next_page_token == "abc" + responses = [] + async for response in async_pager: # pragma: no branch + responses.append(response) + + assert len(responses) == 6 + assert all( + isinstance(i, model_service.ModelVersionCheckpoint) for i in responses + ) + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_async_pages(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), + "__call__", + new_callable=mock.AsyncMock, + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + pages = [] + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch + await client.list_model_version_checkpoints(request={}) + ).pages: + pages.append(page_) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.UpdateModelRequest, + dict, + ], +) +def test_update_model(request_type, transport: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model( + name="name_value", + version_id="version_id_value", + version_aliases=["version_aliases_value"], + display_name="display_name_value", + description="description_value", + version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", + metadata_schema_uri="metadata_schema_uri_value", + training_pipeline="training_pipeline_value", + pipeline_job="pipeline_job_value", + artifact_uri="artifact_uri_value", + supported_deployment_resources_types=[ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ], + supported_input_storage_formats=["supported_input_storage_formats_value"], + supported_output_storage_formats=["supported_output_storage_formats_value"], + etag="etag_value", + metadata_artifact="metadata_artifact_value", + satisfies_pzs=True, + satisfies_pzi=True, + ) + response = client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = model_service.UpdateModelRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, gca_model.Model) + assert response.name == "name_value" + assert response.version_id == "version_id_value" + assert response.version_aliases == ["version_aliases_value"] + assert response.display_name == "display_name_value" + assert response.description == "description_value" + assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" + assert response.metadata_schema_uri == "metadata_schema_uri_value" + assert response.training_pipeline == "training_pipeline_value" + assert response.pipeline_job == "pipeline_job_value" + assert response.artifact_uri == "artifact_uri_value" + assert response.supported_deployment_resources_types == [ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ] + assert response.supported_input_storage_formats == [ + "supported_input_storage_formats_value" + ] + assert response.supported_output_storage_formats == [ + "supported_output_storage_formats_value" + ] + assert response.etag == "etag_value" + assert response.metadata_artifact == "metadata_artifact_value" + assert response.satisfies_pzs is True + assert response.satisfies_pzi is True + + +def test_update_model_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are + # automatically populated, according to AIP-4235, with non-empty requests. + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Populate all string fields in the request which are not UUID4 + # since we want to check that UUID4 are populated automatically + # if they meet the requirements of AIP 4235. + request = model_service.UpdateModelRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client.update_model(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == model_service.UpdateModelRequest() + + +def test_update_model_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.update_model in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[client._transport.update_model] = mock_rpc + request = {} + client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + client.update_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_update_model_async_use_cached_wrapped_rpc( + transport: str = "grpc_asyncio", +): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method_async.wrap_method") as wrapper_fn: + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert ( + client._client._transport.update_model + in client._client._transport._wrapped_methods + ) + + # Replace cached wrapped function with mock + mock_rpc = mock.AsyncMock() + mock_rpc.return_value = mock.Mock() + client._client._transport._wrapped_methods[ + client._client._transport.update_model + ] = mock_rpc + + request = {} + await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + await client.update_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_update_model_async( + transport: str = "grpc_asyncio", request_type=model_service.UpdateModelRequest +): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + gca_model.Model( + name="name_value", + version_id="version_id_value", + version_aliases=["version_aliases_value"], + display_name="display_name_value", + description="description_value", + version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", + metadata_schema_uri="metadata_schema_uri_value", + training_pipeline="training_pipeline_value", + pipeline_job="pipeline_job_value", + artifact_uri="artifact_uri_value", + supported_deployment_resources_types=[ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ], + supported_input_storage_formats=[ + "supported_input_storage_formats_value" + ], + supported_output_storage_formats=[ + "supported_output_storage_formats_value" + ], + etag="etag_value", + metadata_artifact="metadata_artifact_value", + satisfies_pzs=True, + satisfies_pzi=True, + ) + ) + response = await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = model_service.UpdateModelRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, gca_model.Model) + assert response.name == "name_value" + assert response.version_id == "version_id_value" + assert response.version_aliases == ["version_aliases_value"] + assert response.display_name == "display_name_value" + assert response.description == "description_value" + assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" + assert response.metadata_schema_uri == "metadata_schema_uri_value" + assert response.training_pipeline == "training_pipeline_value" + assert response.pipeline_job == "pipeline_job_value" + assert response.artifact_uri == "artifact_uri_value" + assert response.supported_deployment_resources_types == [ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ] + assert response.supported_input_storage_formats == [ + "supported_input_storage_formats_value" + ] + assert response.supported_output_storage_formats == [ + "supported_output_storage_formats_value" + ] + assert response.etag == "etag_value" + assert response.metadata_artifact == "metadata_artifact_value" + assert response.satisfies_pzs is True + assert response.satisfies_pzi is True + + +@pytest.mark.asyncio +async def test_update_model_async_from_dict(): + await test_update_model_async(request_type=dict) + + +def test_update_model_field_headers(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = model_service.UpdateModelRequest() + + request.model.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value = gca_model.Model() + client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "model.name=name_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_update_model_field_headers_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = model_service.UpdateModelRequest() + + request.model.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "model.name=name_value", + ) in kw["metadata"] + + +def test_update_model_flattened(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model() + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.update_model( + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + arg = args[0].model + mock_val = gca_model.Model(name="name_value") + assert arg == mock_val + arg = args[0].update_mask + mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + assert arg == mock_val + + +def test_update_model_flattened_error(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.update_model( + model_service.UpdateModelRequest(), + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + +@pytest.mark.asyncio +async def test_update_model_flattened_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.update_model( + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + arg = args[0].model + mock_val = gca_model.Model(name="name_value") + assert arg == mock_val + arg = args[0].update_mask + mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + assert arg == mock_val + + +@pytest.mark.asyncio +async def test_update_model_flattened_error_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.update_model( + model_service.UpdateModelRequest(), + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.UpdateExplanationDatasetRequest, + dict, + ], +) +def test_update_explanation_dataset(request_type, transport: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.update_explanation_dataset), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/spam") + response = client.update_explanation_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = model_service.UpdateExplanationDatasetRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +def test_update_explanation_dataset_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are # automatically populated, according to AIP-4235, with non-empty requests. client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -4349,6 +4959,7 @@ def test_merge_version_aliases(request_type, transport: str = "grpc"): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -4379,6 +4990,7 @@ def test_merge_version_aliases(request_type, transport: str = "grpc"): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -4537,6 +5149,7 @@ async def test_merge_version_aliases_async( display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -4572,6 +5185,7 @@ async def test_merge_version_aliases_async( assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -8407,20 +9021,202 @@ def test_upload_model_rest_use_cached_wrapped_rpc(): # subsequent calls should use the cached wrapper wrapper_fn.reset_mock() - client.upload_model(request) + client.upload_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +def test_upload_model_rest_required_fields( + request_type=model_service.UploadModelRequest, +): + transport_class = transports.ModelServiceRestTransport + + request_init = {} + request_init["parent"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson(pb_request, use_integers_for_enums=False) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).upload_model._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["parent"] = "parent_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).upload_model._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "parent" in jsonified_request + assert jsonified_request["parent"] == "parent_value" + + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "post", + "query_params": pb_request, + } + transcode_result["body"] = pb_request + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + + response = client.upload_model(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_upload_model_rest_unset_required_fields(): + transport = transports.ModelServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.upload_model._get_unset_required_fields({}) + assert set(unset_fields) == ( + set(()) + & set( + ( + "parent", + "model", + ) + ) + ) + + +def test_upload_model_rest_flattened(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + + # get arguments that satisfy an http rule for this method + sample_request = {"parent": "projects/sample1/locations/sample2"} + + # get truthy value for each flattened field + mock_args = dict( + parent="parent_value", + model=gca_model.Model(name="name_value"), + ) + mock_args.update(sample_request) + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + + client.upload_model(**mock_args) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(req.mock_calls) == 1 + _, args, _ = req.mock_calls[0] + assert path_template.validate( + "%s/v1/{parent=projects/*/locations/*}/models:upload" + % client.transport._host, + args[1], + ) + + +def test_upload_model_rest_flattened_error(transport: str = "rest"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.upload_model( + model_service.UploadModelRequest(), + parent="parent_value", + model=gca_model.Model(name="name_value"), + ) + + +def test_get_model_rest_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.get_model in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[client._transport.get_model] = mock_rpc + + request = {} + client.get_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + client.get_model(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_upload_model_rest_required_fields( - request_type=model_service.UploadModelRequest, -): +def test_get_model_rest_required_fields(request_type=model_service.GetModelRequest): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["parent"] = "" + request_init["name"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8431,21 +9227,21 @@ def test_upload_model_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).upload_model._get_unset_required_fields(jsonified_request) + ).get_model._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["parent"] = "parent_value" + jsonified_request["name"] = "name_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).upload_model._get_unset_required_fields(jsonified_request) + ).get_model._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "parent" in jsonified_request - assert jsonified_request["parent"] == "parent_value" + assert "name" in jsonified_request + assert jsonified_request["name"] == "name_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8454,7 +9250,7 @@ def test_upload_model_rest_required_fields( request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = operations_pb2.Operation(name="operations/spam") + return_value = model.Model() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8466,45 +9262,39 @@ def test_upload_model_rest_required_fields( pb_request = request_type.pb(request) transcode_result = { "uri": "v1/sample_method", - "method": "post", + "method": "get", "query_params": pb_request, } - transcode_result["body"] = pb_request transcode.return_value = transcode_result response_value = Response() response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model.Model.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.upload_model(request) + response = client.get_model(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_upload_model_rest_unset_required_fields(): +def test_get_model_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.upload_model._get_unset_required_fields({}) - assert set(unset_fields) == ( - set(()) - & set( - ( - "parent", - "model", - ) - ) - ) + unset_fields = transport.get_model._get_unset_required_fields({}) + assert set(unset_fields) == (set(()) & set(("name",))) -def test_upload_model_rest_flattened(): +def test_get_model_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8513,40 +9303,40 @@ def test_upload_model_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = operations_pb2.Operation(name="operations/spam") + return_value = model.Model() # get arguments that satisfy an http rule for this method - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} # get truthy value for each flattened field mock_args = dict( - parent="parent_value", - model=gca_model.Model(name="name_value"), + name="name_value", ) mock_args.update(sample_request) # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 + # Convert return value to protobuf type + return_value = model.Model.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.upload_model(**mock_args) + client.get_model(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1/{parent=projects/*/locations/*}/models:upload" - % client.transport._host, + "%s/v1/{name=projects/*/locations/*/models/*}" % client.transport._host, args[1], ) -def test_upload_model_rest_flattened_error(transport: str = "rest"): +def test_get_model_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8555,14 +9345,13 @@ def test_upload_model_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.upload_model( - model_service.UploadModelRequest(), - parent="parent_value", - model=gca_model.Model(name="name_value"), + client.get_model( + model_service.GetModelRequest(), + name="name_value", ) -def test_get_model_rest_use_cached_wrapped_rpc(): +def test_list_models_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -8576,33 +9365,33 @@ def test_get_model_rest_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.get_model in client._transport._wrapped_methods + assert client._transport.list_models in client._transport._wrapped_methods # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.get_model] = mock_rpc + client._transport._wrapped_methods[client._transport.list_models] = mock_rpc request = {} - client.get_model(request) + client.list_models(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.get_model(request) + client.list_models(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_get_model_rest_required_fields(request_type=model_service.GetModelRequest): +def test_list_models_rest_required_fields(request_type=model_service.ListModelsRequest): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["name"] = "" + request_init["parent"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8613,21 +9402,31 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).get_model._get_unset_required_fields(jsonified_request) + ).list_models._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["name"] = "name_value" + jsonified_request["parent"] = "parent_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).get_model._get_unset_required_fields(jsonified_request) + ).list_models._get_unset_required_fields(jsonified_request) + # Check that path parameters and body parameters are not mixing in. + assert not set(unset_fields) - set( + ( + "filter", + "order_by", + "page_size", + "page_token", + "read_mask", + ) + ) jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "name" in jsonified_request - assert jsonified_request["name"] == "name_value" + assert "parent" in jsonified_request + assert jsonified_request["parent"] == "parent_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8636,7 +9435,7 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model.Model() + return_value = model_service.ListModelsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8657,30 +9456,41 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque response_value.status_code = 200 # Convert return value to protobuf type - return_value = model.Model.pb(return_value) + return_value = model_service.ListModelsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.get_model(request) + response = client.list_models(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_get_model_rest_unset_required_fields(): +def test_list_models_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.get_model._get_unset_required_fields({}) - assert set(unset_fields) == (set(()) & set(("name",))) + unset_fields = transport.list_models._get_unset_required_fields({}) + assert set(unset_fields) == ( + set( + ( + "filter", + "orderBy", + "pageSize", + "pageToken", + "readMask", + ) + ) + & set(("parent",)) + ) -def test_get_model_rest_flattened(): +def test_list_models_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8689,14 +9499,14 @@ def test_get_model_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model.Model() + return_value = model_service.ListModelsResponse() # get arguments that satisfy an http rule for this method - sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} + sample_request = {"parent": "projects/sample1/locations/sample2"} # get truthy value for each flattened field mock_args = dict( - name="name_value", + parent="parent_value", ) mock_args.update(sample_request) @@ -8704,25 +9514,25 @@ def test_get_model_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model.Model.pb(return_value) + return_value = model_service.ListModelsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.get_model(**mock_args) + client.list_models(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1/{name=projects/*/locations/*/models/*}" % client.transport._host, + "%s/v1/{parent=projects/*/locations/*}/models" % client.transport._host, args[1], ) -def test_get_model_rest_flattened_error(transport: str = "rest"): +def test_list_models_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8731,13 +9541,74 @@ def test_get_model_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.get_model( - model_service.GetModelRequest(), - name="name_value", + client.list_models( + model_service.ListModelsRequest(), + parent="parent_value", ) -def test_list_models_rest_use_cached_wrapped_rpc(): +def test_list_models_rest_pager(transport: str = "rest"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # TODO(kbandes): remove this mock unless there's a good reason for it. + # with mock.patch.object(path_template, 'transcode') as transcode: + # Set the response as a series of pages + response = ( + model_service.ListModelsResponse( + models=[ + model.Model(), + model.Model(), + model.Model(), + ], + next_page_token="abc", + ), + model_service.ListModelsResponse( + models=[], + next_page_token="def", + ), + model_service.ListModelsResponse( + models=[ + model.Model(), + ], + next_page_token="ghi", + ), + model_service.ListModelsResponse( + models=[ + model.Model(), + model.Model(), + ], + ), + ) + # Two responses for two calls + response = response + response + + # Wrap the values into proper Response objs + response = tuple(model_service.ListModelsResponse.to_json(x) for x in response) + return_values = tuple(Response() for i in response) + for return_val, response_val in zip(return_values, response): + return_val._content = response_val.encode("UTF-8") + return_val.status_code = 200 + req.side_effect = return_values + + sample_request = {"parent": "projects/sample1/locations/sample2"} + + pager = client.list_models(request=sample_request) + + results = list(pager) + assert len(results) == 6 + assert all(isinstance(i, model.Model) for i in results) + + pages = list(client.list_models(request=sample_request).pages) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +def test_list_model_versions_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -8751,33 +9622,39 @@ def test_list_models_rest_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.list_models in client._transport._wrapped_methods + assert ( + client._transport.list_model_versions in client._transport._wrapped_methods + ) # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.list_models] = mock_rpc + client._transport._wrapped_methods[ + client._transport.list_model_versions + ] = mock_rpc request = {} - client.list_models(request) + client.list_model_versions(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.list_models(request) + client.list_model_versions(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_list_models_rest_required_fields(request_type=model_service.ListModelsRequest): +def test_list_model_versions_rest_required_fields( + request_type=model_service.ListModelVersionsRequest, +): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["parent"] = "" + request_init["name"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8788,16 +9665,16 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_models._get_unset_required_fields(jsonified_request) + ).list_model_versions._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["parent"] = "parent_value" + jsonified_request["name"] = "name_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_models._get_unset_required_fields(jsonified_request) + ).list_model_versions._get_unset_required_fields(jsonified_request) # Check that path parameters and body parameters are not mixing in. assert not set(unset_fields) - set( ( @@ -8811,8 +9688,8 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "parent" in jsonified_request - assert jsonified_request["parent"] == "parent_value" + assert "name" in jsonified_request + assert jsonified_request["name"] == "name_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8821,7 +9698,7 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8842,26 +9719,26 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelsResponse.pb(return_value) + return_value = model_service.ListModelVersionsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_models(request) + response = client.list_model_versions(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_list_models_rest_unset_required_fields(): +def test_list_model_versions_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.list_models._get_unset_required_fields({}) + unset_fields = transport.list_model_versions._get_unset_required_fields({}) assert set(unset_fields) == ( set( ( @@ -8872,11 +9749,11 @@ def test_list_models_rest_unset_required_fields(): "readMask", ) ) - & set(("parent",)) + & set(("name",)) ) -def test_list_models_rest_flattened(): +def test_list_model_versions_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8885,14 +9762,14 @@ def test_list_models_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse() # get arguments that satisfy an http rule for this method - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} # get truthy value for each flattened field mock_args = dict( - parent="parent_value", + name="name_value", ) mock_args.update(sample_request) @@ -8900,25 +9777,26 @@ def test_list_models_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelsResponse.pb(return_value) + return_value = model_service.ListModelVersionsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_models(**mock_args) + client.list_model_versions(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1/{parent=projects/*/locations/*}/models" % client.transport._host, + "%s/v1/{name=projects/*/locations/*/models/*}:listVersions" + % client.transport._host, args[1], ) -def test_list_models_rest_flattened_error(transport: str = "rest"): +def test_list_model_versions_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8927,13 +9805,13 @@ def test_list_models_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.list_models( - model_service.ListModelsRequest(), - parent="parent_value", + client.list_model_versions( + model_service.ListModelVersionsRequest(), + name="name_value", ) -def test_list_models_rest_pager(transport: str = "rest"): +def test_list_model_versions_rest_pager(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8945,7 +9823,7 @@ def test_list_models_rest_pager(transport: str = "rest"): # with mock.patch.object(path_template, 'transcode') as transcode: # Set the response as a series of pages response = ( - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), model.Model(), @@ -8953,17 +9831,17 @@ def test_list_models_rest_pager(transport: str = "rest"): ], next_page_token="abc", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[], next_page_token="def", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), ], next_page_token="ghi", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), model.Model(), @@ -8974,27 +9852,29 @@ def test_list_models_rest_pager(transport: str = "rest"): response = response + response # Wrap the values into proper Response objs - response = tuple(model_service.ListModelsResponse.to_json(x) for x in response) + response = tuple( + model_service.ListModelVersionsResponse.to_json(x) for x in response + ) return_values = tuple(Response() for i in response) for return_val, response_val in zip(return_values, response): return_val._content = response_val.encode("UTF-8") return_val.status_code = 200 req.side_effect = return_values - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} - pager = client.list_models(request=sample_request) + pager = client.list_model_versions(request=sample_request) results = list(pager) assert len(results) == 6 assert all(isinstance(i, model.Model) for i in results) - pages = list(client.list_models(request=sample_request).pages) + pages = list(client.list_model_versions(request=sample_request).pages) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token -def test_list_model_versions_rest_use_cached_wrapped_rpc(): +def test_list_model_version_checkpoints_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -9009,7 +9889,8 @@ def test_list_model_versions_rest_use_cached_wrapped_rpc(): # Ensure method has been cached assert ( - client._transport.list_model_versions in client._transport._wrapped_methods + client._transport.list_model_version_checkpoints + in client._transport._wrapped_methods ) # Replace cached wrapped function with mock @@ -9018,24 +9899,24 @@ def test_list_model_versions_rest_use_cached_wrapped_rpc(): "foo" # operation_request.operation in compute client(s) expect a string. ) client._transport._wrapped_methods[ - client._transport.list_model_versions + client._transport.list_model_version_checkpoints ] = mock_rpc request = {} - client.list_model_versions(request) + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.list_model_versions(request) + client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_list_model_versions_rest_required_fields( - request_type=model_service.ListModelVersionsRequest, +def test_list_model_version_checkpoints_rest_required_fields( + request_type=model_service.ListModelVersionCheckpointsRequest, ): transport_class = transports.ModelServiceRestTransport @@ -9051,7 +9932,7 @@ def test_list_model_versions_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_model_versions._get_unset_required_fields(jsonified_request) + ).list_model_version_checkpoints._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present @@ -9060,15 +9941,12 @@ def test_list_model_versions_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_model_versions._get_unset_required_fields(jsonified_request) + ).list_model_version_checkpoints._get_unset_required_fields(jsonified_request) # Check that path parameters and body parameters are not mixing in. assert not set(unset_fields) - set( ( - "filter", - "order_by", "page_size", "page_token", - "read_mask", ) ) jsonified_request.update(unset_fields) @@ -9084,7 +9962,7 @@ def test_list_model_versions_rest_required_fields( request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -9105,41 +9983,42 @@ def test_list_model_versions_rest_required_fields( response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_model_versions(request) + response = client.list_model_version_checkpoints(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_list_model_versions_rest_unset_required_fields(): +def test_list_model_version_checkpoints_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.list_model_versions._get_unset_required_fields({}) + unset_fields = transport.list_model_version_checkpoints._get_unset_required_fields( + {} + ) assert set(unset_fields) == ( set( ( - "filter", - "orderBy", "pageSize", "pageToken", - "readMask", ) ) & set(("name",)) ) -def test_list_model_versions_rest_flattened(): +def test_list_model_version_checkpoints_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -9148,7 +10027,7 @@ def test_list_model_versions_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse() # get arguments that satisfy an http rule for this method sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} @@ -9163,26 +10042,28 @@ def test_list_model_versions_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_model_versions(**mock_args) + client.list_model_version_checkpoints(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1/{name=projects/*/locations/*/models/*}:listVersions" + "%s/v1/{name=projects/*/locations/*/models/*}:listCheckpoints" % client.transport._host, args[1], ) -def test_list_model_versions_rest_flattened_error(transport: str = "rest"): +def test_list_model_version_checkpoints_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -9191,13 +10072,13 @@ def test_list_model_versions_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.list_model_versions( - model_service.ListModelVersionsRequest(), + client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), name="name_value", ) -def test_list_model_versions_rest_pager(transport: str = "rest"): +def test_list_model_version_checkpoints_rest_pager(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -9209,28 +10090,28 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): # with mock.patch.object(path_template, 'transcode') as transcode: # Set the response as a series of pages response = ( - model_service.ListModelVersionsResponse( - models=[ - model.Model(), - model.Model(), - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), ], next_page_token="abc", ), - model_service.ListModelVersionsResponse( - models=[], + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], next_page_token="def", ), - model_service.ListModelVersionsResponse( - models=[ - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), ], next_page_token="ghi", ), - model_service.ListModelVersionsResponse( - models=[ - model.Model(), - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), ], ), ) @@ -9239,7 +10120,8 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): # Wrap the values into proper Response objs response = tuple( - model_service.ListModelVersionsResponse.to_json(x) for x in response + model_service.ListModelVersionCheckpointsResponse.to_json(x) + for x in response ) return_values = tuple(Response() for i in response) for return_val, response_val in zip(return_values, response): @@ -9249,13 +10131,15 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} - pager = client.list_model_versions(request=sample_request) + pager = client.list_model_version_checkpoints(request=sample_request) results = list(pager) assert len(results) == 6 - assert all(isinstance(i, model.Model) for i in results) + assert all(isinstance(i, model_service.ModelVersionCheckpoint) for i in results) - pages = list(client.list_model_versions(request=sample_request).pages) + pages = list( + client.list_model_version_checkpoints(request=sample_request).pages + ) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -12281,6 +13165,29 @@ def test_list_model_versions_empty_call_grpc(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_list_model_version_checkpoints_empty_call_grpc(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = model_service.ListModelVersionCheckpointsResponse() + client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. def test_update_model_empty_call_grpc(): @@ -12654,6 +13561,7 @@ async def test_get_model_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -12739,6 +13647,35 @@ async def test_list_model_versions_empty_call_grpc_asyncio(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_empty_call_grpc_asyncio(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", + ) + ) + await client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @pytest.mark.asyncio @@ -12759,6 +13696,7 @@ async def test_update_model_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -12889,6 +13827,7 @@ async def test_merge_version_aliases_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -13267,10 +14206,13 @@ def test_upload_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_upload_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_upload_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_upload_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UploadModelRequest.pb( model_service.UploadModelRequest() ) @@ -13294,6 +14236,7 @@ def test_upload_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.upload_model( request, @@ -13305,6 +14248,7 @@ def test_upload_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_rest_bad_request(request_type=model_service.GetModelRequest): @@ -13356,6 +14300,7 @@ def test_get_model_rest_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -13391,6 +14336,7 @@ def test_get_model_rest_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -13427,10 +14373,13 @@ def test_get_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_get_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelRequest.pb(model_service.GetModelRequest()) transcode.return_value = { "method": "post", @@ -13452,6 +14401,7 @@ def test_get_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata client.get_model( request, @@ -13463,6 +14413,7 @@ def test_get_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_models_rest_bad_request(request_type=model_service.ListModelsRequest): @@ -13521,15 +14472,144 @@ def test_list_models_rest_call_success(request_type): response_value.content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_models(request) + response = client.list_models(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, pagers.ListModelsPager) + assert response.next_page_token == "next_page_token_value" + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_list_models_rest_interceptors(null_interceptor): + transport = transports.ModelServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.ModelServiceRestInterceptor(), + ) + client = ModelServiceClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_list_models" + ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_list_models_with_metadata" + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_models" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelsRequest.pb( + model_service.ListModelsRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = model_service.ListModelsResponse.to_json( + model_service.ListModelsResponse() + ) + req.return_value.content = return_value + + request = model_service.ListModelsRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = model_service.ListModelsResponse() + post_with_metadata.return_value = model_service.ListModelsResponse(), metadata + + client.list_models( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() + + +def test_list_model_versions_rest_bad_request( + request_type=model_service.ListModelVersionsRequest, +): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + json_return_value = "" + response_value.json = mock.Mock(return_value={}) + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + client.list_model_versions(request) + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.ListModelVersionsRequest, + dict, + ], +) +def test_list_model_versions_rest_call_success(request_type): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = model_service.ListModelVersionsResponse( + next_page_token="next_page_token_value", + ) + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model_service.ListModelVersionsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + response_value.content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = client.list_model_versions(request) # Establish that the response is the type that we expect. - assert isinstance(response, pagers.ListModelsPager) + assert isinstance(response, pagers.ListModelVersionsPager) assert response.next_page_token == "next_page_token_value" @pytest.mark.parametrize("null_interceptor", [True, False]) -def test_list_models_rest_interceptors(null_interceptor): +def test_list_model_versions_rest_interceptors(null_interceptor): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials(), interceptor=None @@ -13543,14 +14623,17 @@ def test_list_models_rest_interceptors(null_interceptor): ) as req, mock.patch.object( path_template, "transcode" ) as transcode, mock.patch.object( - transports.ModelServiceRestInterceptor, "post_list_models" + transports.ModelServiceRestInterceptor, "post_list_model_versions" ) as post, mock.patch.object( - transports.ModelServiceRestInterceptor, "pre_list_models" + transports.ModelServiceRestInterceptor, "post_list_model_versions_with_metadata" + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_model_versions" ) as pre: pre.assert_not_called() post.assert_not_called() - pb_message = model_service.ListModelsRequest.pb( - model_service.ListModelsRequest() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionsRequest.pb( + model_service.ListModelVersionsRequest() ) transcode.return_value = { "method": "post", @@ -13562,20 +14645,24 @@ def test_list_models_rest_interceptors(null_interceptor): req.return_value = mock.Mock() req.return_value.status_code = 200 req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - return_value = model_service.ListModelsResponse.to_json( - model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse.to_json( + model_service.ListModelVersionsResponse() ) req.return_value.content = return_value - request = model_service.ListModelsRequest() + request = model_service.ListModelVersionsRequest() metadata = [ ("key", "val"), ("cephalopod", "squid"), ] pre.return_value = request, metadata - post.return_value = model_service.ListModelsResponse() + post.return_value = model_service.ListModelVersionsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionsResponse(), + metadata, + ) - client.list_models( + client.list_model_versions( request, metadata=[ ("key", "val"), @@ -13585,10 +14672,11 @@ def test_list_models_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() -def test_list_model_versions_rest_bad_request( - request_type=model_service.ListModelVersionsRequest, +def test_list_model_version_checkpoints_rest_bad_request( + request_type=model_service.ListModelVersionCheckpointsRequest, ): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" @@ -13609,17 +14697,17 @@ def test_list_model_versions_rest_bad_request( response_value.request = mock.Mock() req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_model_versions(request) + client.list_model_version_checkpoints(request) @pytest.mark.parametrize( "request_type", [ - model_service.ListModelVersionsRequest, + model_service.ListModelVersionCheckpointsRequest, dict, ], ) -def test_list_model_versions_rest_call_success(request_type): +def test_list_model_version_checkpoints_rest_call_success(request_type): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" ) @@ -13631,7 +14719,7 @@ def test_list_model_versions_rest_call_success(request_type): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse( + return_value = model_service.ListModelVersionCheckpointsResponse( next_page_token="next_page_token_value", ) @@ -13640,20 +14728,22 @@ def test_list_model_versions_rest_call_success(request_type): response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value.content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_model_versions(request) + response = client.list_model_version_checkpoints(request) # Establish that the response is the type that we expect. - assert isinstance(response, pagers.ListModelVersionsPager) + assert isinstance(response, pagers.ListModelVersionCheckpointsPager) assert response.next_page_token == "next_page_token_value" @pytest.mark.parametrize("null_interceptor", [True, False]) -def test_list_model_versions_rest_interceptors(null_interceptor): +def test_list_model_version_checkpoints_rest_interceptors(null_interceptor): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials(), interceptor=None @@ -13667,14 +14757,18 @@ def test_list_model_versions_rest_interceptors(null_interceptor): ) as req, mock.patch.object( path_template, "transcode" ) as transcode, mock.patch.object( - transports.ModelServiceRestInterceptor, "post_list_model_versions" + transports.ModelServiceRestInterceptor, "post_list_model_version_checkpoints" ) as post, mock.patch.object( - transports.ModelServiceRestInterceptor, "pre_list_model_versions" + transports.ModelServiceRestInterceptor, + "post_list_model_version_checkpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_model_version_checkpoints" ) as pre: pre.assert_not_called() post.assert_not_called() - pb_message = model_service.ListModelVersionsRequest.pb( - model_service.ListModelVersionsRequest() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionCheckpointsRequest.pb( + model_service.ListModelVersionCheckpointsRequest() ) transcode.return_value = { "method": "post", @@ -13686,20 +14780,24 @@ def test_list_model_versions_rest_interceptors(null_interceptor): req.return_value = mock.Mock() req.return_value.status_code = 200 req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - return_value = model_service.ListModelVersionsResponse.to_json( - model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse.to_json( + model_service.ListModelVersionCheckpointsResponse() ) req.return_value.content = return_value - request = model_service.ListModelVersionsRequest() + request = model_service.ListModelVersionCheckpointsRequest() metadata = [ ("key", "val"), ("cephalopod", "squid"), ] pre.return_value = request, metadata - post.return_value = model_service.ListModelVersionsResponse() + post.return_value = model_service.ListModelVersionCheckpointsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionCheckpointsResponse(), + metadata, + ) - client.list_model_versions( + client.list_model_version_checkpoints( request, metadata=[ ("key", "val"), @@ -13709,6 +14807,7 @@ def test_list_model_versions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_model_rest_bad_request(request_type=model_service.UpdateModelRequest): @@ -13761,6 +14860,7 @@ def test_update_model_rest_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -13791,10 +14891,23 @@ def test_update_model_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [{"name": "name_value", "value": "value_value"}], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -13952,6 +15065,7 @@ def get_message_fields(field): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -13987,6 +15101,7 @@ def get_message_fields(field): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -14023,10 +15138,13 @@ def test_update_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_update_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_update_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_update_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateModelRequest.pb( model_service.UpdateModelRequest() ) @@ -14050,6 +15168,7 @@ def test_update_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model.Model() + post_with_metadata.return_value = gca_model.Model(), metadata client.update_model( request, @@ -14061,6 +15180,7 @@ def test_update_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_explanation_dataset_rest_bad_request( @@ -14141,10 +15261,14 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_update_explanation_dataset" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_update_explanation_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_update_explanation_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateExplanationDatasetRequest.pb( model_service.UpdateExplanationDatasetRequest() ) @@ -14168,6 +15292,7 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_explanation_dataset( request, @@ -14179,6 +15304,7 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_rest_bad_request(request_type=model_service.DeleteModelRequest): @@ -14257,10 +15383,13 @@ def test_delete_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_delete_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_delete_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_delete_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelRequest.pb( model_service.DeleteModelRequest() ) @@ -14284,6 +15413,7 @@ def test_delete_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model( request, @@ -14295,6 +15425,7 @@ def test_delete_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_version_rest_bad_request( @@ -14375,10 +15506,14 @@ def test_delete_model_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_delete_model_version" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_delete_model_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_delete_model_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelVersionRequest.pb( model_service.DeleteModelVersionRequest() ) @@ -14402,6 +15537,7 @@ def test_delete_model_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_version( request, @@ -14413,6 +15549,7 @@ def test_delete_model_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_merge_version_aliases_rest_bad_request( @@ -14466,6 +15603,7 @@ def test_merge_version_aliases_rest_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -14501,6 +15639,7 @@ def test_merge_version_aliases_rest_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -14537,10 +15676,14 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_merge_version_aliases" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_merge_version_aliases_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_merge_version_aliases" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.MergeVersionAliasesRequest.pb( model_service.MergeVersionAliasesRequest() ) @@ -14564,6 +15707,7 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata client.merge_version_aliases( request, @@ -14575,6 +15719,7 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_model_rest_bad_request(request_type=model_service.ExportModelRequest): @@ -14653,10 +15798,13 @@ def test_export_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_export_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_export_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_export_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ExportModelRequest.pb( model_service.ExportModelRequest() ) @@ -14680,6 +15828,7 @@ def test_export_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_model( request, @@ -14691,6 +15840,7 @@ def test_export_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_copy_model_rest_bad_request(request_type=model_service.CopyModelRequest): @@ -14769,10 +15919,13 @@ def test_copy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_copy_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_copy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_copy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.CopyModelRequest.pb(model_service.CopyModelRequest()) transcode.return_value = { "method": "post", @@ -14794,6 +15947,7 @@ def test_copy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.copy_model( request, @@ -14805,6 +15959,7 @@ def test_copy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_model_evaluation_rest_bad_request( @@ -14899,10 +16054,14 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_import_model_evaluation" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_import_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_import_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ImportModelEvaluationRequest.pb( model_service.ImportModelEvaluationRequest() ) @@ -14928,6 +16087,10 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model_evaluation.ModelEvaluation() + post_with_metadata.return_value = ( + gca_model_evaluation.ModelEvaluation(), + metadata, + ) client.import_model_evaluation( request, @@ -14939,6 +16102,7 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_import_model_evaluation_slices_rest_bad_request( @@ -15032,11 +16196,15 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor transports.ModelServiceRestInterceptor, "post_batch_import_model_evaluation_slices", ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_batch_import_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_batch_import_model_evaluation_slices", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportModelEvaluationSlicesRequest.pb( model_service.BatchImportModelEvaluationSlicesRequest() ) @@ -15062,6 +16230,10 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = model_service.BatchImportModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.BatchImportModelEvaluationSlicesResponse(), + metadata, + ) client.batch_import_model_evaluation_slices( request, @@ -15073,6 +16245,7 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_import_evaluated_annotations_rest_bad_request( @@ -15164,10 +16337,14 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): transports.ModelServiceRestInterceptor, "post_batch_import_evaluated_annotations", ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_batch_import_evaluated_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_batch_import_evaluated_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportEvaluatedAnnotationsRequest.pb( model_service.BatchImportEvaluatedAnnotationsRequest() ) @@ -15193,6 +16370,10 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.BatchImportEvaluatedAnnotationsResponse() + post_with_metadata.return_value = ( + model_service.BatchImportEvaluatedAnnotationsResponse(), + metadata, + ) client.batch_import_evaluated_annotations( request, @@ -15204,6 +16385,7 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_evaluation_rest_bad_request( @@ -15302,10 +16484,14 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model_evaluation" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_get_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationRequest.pb( model_service.GetModelEvaluationRequest() ) @@ -15331,6 +16517,7 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation.ModelEvaluation() + post_with_metadata.return_value = model_evaluation.ModelEvaluation(), metadata client.get_model_evaluation( request, @@ -15342,6 +16529,7 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_evaluations_rest_bad_request( @@ -15426,10 +16614,14 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_list_model_evaluations" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_list_model_evaluations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_list_model_evaluations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationsRequest.pb( model_service.ListModelEvaluationsRequest() ) @@ -15455,6 +16647,10 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationsResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationsResponse(), + metadata, + ) client.list_model_evaluations( request, @@ -15466,6 +16662,7 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_evaluation_slice_rest_bad_request( @@ -15556,10 +16753,14 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model_evaluation_slice" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_get_model_evaluation_slice_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model_evaluation_slice" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationSliceRequest.pb( model_service.GetModelEvaluationSliceRequest() ) @@ -15585,6 +16786,10 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation_slice.ModelEvaluationSlice() + post_with_metadata.return_value = ( + model_evaluation_slice.ModelEvaluationSlice(), + metadata, + ) client.get_model_evaluation_slice( request, @@ -15596,6 +16801,7 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_evaluation_slices_rest_bad_request( @@ -15684,10 +16890,14 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_list_model_evaluation_slices" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_list_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_list_model_evaluation_slices" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationSlicesRequest.pb( model_service.ListModelEvaluationSlicesRequest() ) @@ -15713,6 +16923,10 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationSlicesResponse(), + metadata, + ) client.list_model_evaluation_slices( request, @@ -15724,6 +16938,7 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -16440,6 +17655,28 @@ def test_list_model_versions_empty_call_rest(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_list_model_version_checkpoints_empty_call_rest(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. def test_update_model_empty_call_rest(): @@ -16862,10 +18099,13 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_upload_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_upload_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_upload_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UploadModelRequest.pb( model_service.UploadModelRequest() ) @@ -16889,6 +18129,7 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.upload_model( request, @@ -16900,6 +18141,7 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16962,6 +18204,7 @@ async def test_get_model_rest_asyncio_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -16999,6 +18242,7 @@ async def test_get_model_rest_asyncio_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -17040,10 +18284,13 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_get_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelRequest.pb(model_service.GetModelRequest()) transcode.return_value = { "method": "post", @@ -17065,6 +18312,7 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata await client.get_model( request, @@ -17076,6 +18324,7 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17176,10 +18425,13 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_models" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_list_models_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_models" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelsRequest.pb( model_service.ListModelsRequest() ) @@ -17205,6 +18457,7 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelsResponse() + post_with_metadata.return_value = model_service.ListModelsResponse(), metadata await client.list_models( request, @@ -17216,6 +18469,7 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17316,10 +18570,14 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_versions" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelVersionsRequest.pb( model_service.ListModelVersionsRequest() ) @@ -17345,6 +18603,10 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelVersionsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionsResponse(), + metadata, + ) await client.list_model_versions( request, @@ -17356,6 +18618,162 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_rest_asyncio_bad_request( + request_type=model_service.ListModelVersionCheckpointsRequest, +): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(AsyncAuthorizedSession, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.read = mock.AsyncMock(return_value=b"{}") + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + await client.list_model_version_checkpoints(request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_type", + [ + model_service.ListModelVersionCheckpointsRequest, + dict, + ], +) +async def test_list_model_version_checkpoints_rest_asyncio_call_success(request_type): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", + ) + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) + json_return_value = json_format.MessageToJson(return_value) + response_value.read = mock.AsyncMock( + return_value=json_return_value.encode("UTF-8") + ) + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = await client.list_model_version_checkpoints(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, pagers.ListModelVersionCheckpointsAsyncPager) + assert response.next_page_token == "next_page_token_value" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("null_interceptor", [True, False]) +async def test_list_model_version_checkpoints_rest_asyncio_interceptors( + null_interceptor, +): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + transport = transports.AsyncModelServiceRestTransport( + credentials=async_anonymous_credentials(), + interceptor=None + if null_interceptor + else transports.AsyncModelServiceRestInterceptor(), + ) + client = ModelServiceAsyncClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_version_checkpoints", + ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_version_checkpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "pre_list_model_version_checkpoints", + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionCheckpointsRequest.pb( + model_service.ListModelVersionCheckpointsRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = model_service.ListModelVersionCheckpointsResponse.to_json( + model_service.ListModelVersionCheckpointsResponse() + ) + req.return_value.read = mock.AsyncMock(return_value=return_value) + + request = model_service.ListModelVersionCheckpointsRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = model_service.ListModelVersionCheckpointsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionCheckpointsResponse(), + metadata, + ) + + await client.list_model_version_checkpoints( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17419,6 +18837,7 @@ async def test_update_model_rest_asyncio_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -17449,10 +18868,23 @@ async def test_update_model_rest_asyncio_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [{"name": "name_value", "value": "value_value"}], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -17610,6 +19042,7 @@ def get_message_fields(field): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -17647,6 +19080,7 @@ def get_message_fields(field): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -17688,10 +19122,13 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_update_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_update_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_update_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateModelRequest.pb( model_service.UpdateModelRequest() ) @@ -17715,6 +19152,7 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model.Model() + post_with_metadata.return_value = gca_model.Model(), metadata await client.update_model( request, @@ -17726,6 +19164,7 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17822,10 +19261,14 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_update_explanation_dataset" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_update_explanation_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_update_explanation_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateExplanationDatasetRequest.pb( model_service.UpdateExplanationDatasetRequest() ) @@ -17849,6 +19292,7 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_explanation_dataset( request, @@ -17860,6 +19304,7 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17956,10 +19401,13 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_delete_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_delete_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_delete_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelRequest.pb( model_service.DeleteModelRequest() ) @@ -17983,6 +19431,7 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model( request, @@ -17994,6 +19443,7 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18090,10 +19540,14 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_delete_model_version" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_delete_model_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_delete_model_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelVersionRequest.pb( model_service.DeleteModelVersionRequest() ) @@ -18117,6 +19571,7 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_version( request, @@ -18128,6 +19583,7 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18190,6 +19646,7 @@ async def test_merge_version_aliases_rest_asyncio_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", pipeline_job="pipeline_job_value", @@ -18227,6 +19684,7 @@ async def test_merge_version_aliases_rest_asyncio_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.pipeline_job == "pipeline_job_value" @@ -18268,10 +19726,14 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_merge_version_aliases" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_merge_version_aliases_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_merge_version_aliases" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.MergeVersionAliasesRequest.pb( model_service.MergeVersionAliasesRequest() ) @@ -18295,6 +19757,7 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata await client.merge_version_aliases( request, @@ -18306,6 +19769,7 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18402,10 +19866,13 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_export_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_export_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_export_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ExportModelRequest.pb( model_service.ExportModelRequest() ) @@ -18429,6 +19896,7 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_model( request, @@ -18440,6 +19908,7 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18536,10 +20005,13 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_copy_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_copy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_copy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.CopyModelRequest.pb(model_service.CopyModelRequest()) transcode.return_value = { "method": "post", @@ -18561,6 +20033,7 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.copy_model( request, @@ -18572,6 +20045,7 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18682,10 +20156,14 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_import_model_evaluation" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_import_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_import_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ImportModelEvaluationRequest.pb( model_service.ImportModelEvaluationRequest() ) @@ -18711,6 +20189,10 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = gca_model_evaluation.ModelEvaluation() + post_with_metadata.return_value = ( + gca_model_evaluation.ModelEvaluation(), + metadata, + ) await client.import_model_evaluation( request, @@ -18722,6 +20204,7 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18835,11 +20318,15 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( transports.AsyncModelServiceRestInterceptor, "post_batch_import_model_evaluation_slices", ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_batch_import_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_batch_import_model_evaluation_slices", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportModelEvaluationSlicesRequest.pb( model_service.BatchImportModelEvaluationSlicesRequest() ) @@ -18865,6 +20352,10 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = model_service.BatchImportModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.BatchImportModelEvaluationSlicesResponse(), + metadata, + ) await client.batch_import_model_evaluation_slices( request, @@ -18876,6 +20367,7 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18987,11 +20479,15 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( transports.AsyncModelServiceRestInterceptor, "post_batch_import_evaluated_annotations", ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_batch_import_evaluated_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_batch_import_evaluated_annotations", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportEvaluatedAnnotationsRequest.pb( model_service.BatchImportEvaluatedAnnotationsRequest() ) @@ -19017,6 +20513,10 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = model_service.BatchImportEvaluatedAnnotationsResponse() + post_with_metadata.return_value = ( + model_service.BatchImportEvaluatedAnnotationsResponse(), + metadata, + ) await client.batch_import_evaluated_annotations( request, @@ -19028,6 +20528,7 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19142,10 +20643,14 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model_evaluation" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_get_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationRequest.pb( model_service.GetModelEvaluationRequest() ) @@ -19171,6 +20676,7 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation.ModelEvaluation() + post_with_metadata.return_value = model_evaluation.ModelEvaluation(), metadata await client.get_model_evaluation( request, @@ -19182,6 +20688,7 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19282,10 +20789,14 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_evaluations" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_evaluations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_evaluations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationsRequest.pb( model_service.ListModelEvaluationsRequest() ) @@ -19311,6 +20822,10 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationsResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationsResponse(), + metadata, + ) await client.list_model_evaluations( request, @@ -19322,6 +20837,7 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19428,10 +20944,14 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model_evaluation_slice" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_get_model_evaluation_slice_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model_evaluation_slice" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationSliceRequest.pb( model_service.GetModelEvaluationSliceRequest() ) @@ -19457,6 +20977,10 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = model_evaluation_slice.ModelEvaluationSlice() + post_with_metadata.return_value = ( + model_evaluation_slice.ModelEvaluationSlice(), + metadata, + ) await client.get_model_evaluation_slice( request, @@ -19468,6 +20992,7 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19572,10 +21097,14 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_evaluation_slices" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_evaluation_slices" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationSlicesRequest.pb( model_service.ListModelEvaluationSlicesRequest() ) @@ -19601,6 +21130,10 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationSlicesResponse(), + metadata, + ) await client.list_model_evaluation_slices( request, @@ -19612,6 +21145,7 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20464,6 +21998,33 @@ async def test_list_model_versions_empty_call_rest_asyncio(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_empty_call_rest_asyncio(): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="rest_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + await client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @pytest.mark.asyncio @@ -20906,6 +22467,7 @@ def test_model_service_base_transport(): "get_model", "list_models", "list_model_versions", + "list_model_version_checkpoints", "update_model", "update_explanation_dataset", "delete_model", @@ -21202,6 +22764,9 @@ def test_model_service_client_transport_session_collision(transport_name): session1 = client1.transport.list_model_versions._session session2 = client2.transport.list_model_versions._session assert session1 != session2 + session1 = client1.transport.list_model_version_checkpoints._session + session2 = client2.transport.list_model_version_checkpoints._session + assert session1 != session2 session1 = client1.transport.update_model._session session2 = client2.transport.update_model._session assert session1 != session2 diff --git a/tests/unit/gapic/aiplatform_v1/test_notebook_service.py b/tests/unit/gapic/aiplatform_v1/test_notebook_service.py index 27ad674bbe..f920057b30 100644 --- a/tests/unit/gapic/aiplatform_v1/test_notebook_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_notebook_service.py @@ -74,6 +74,7 @@ from google.cloud.aiplatform_v1.services.notebook_service import transports from google.cloud.aiplatform_v1.types import accelerator_type from google.cloud.aiplatform_v1.types import encryption_spec +from google.cloud.aiplatform_v1.types import env_var from google.cloud.aiplatform_v1.types import job_state from google.cloud.aiplatform_v1.types import machine_resources from google.cloud.aiplatform_v1.types import network_spec @@ -87,6 +88,7 @@ from google.cloud.aiplatform_v1.types import notebook_runtime as gca_notebook_runtime from google.cloud.aiplatform_v1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1.types import notebook_service +from google.cloud.aiplatform_v1.types import notebook_software_config from google.cloud.aiplatform_v1.types import operation as gca_operation from google.cloud.aiplatform_v1.types import reservation_affinity from google.cloud.location import locations_pb2 @@ -104,6 +106,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -361,6 +371,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = NotebookServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = NotebookServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -11916,6 +11969,14 @@ def test_create_notebook_runtime_template_rest_call_success(request_type): "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -12030,11 +12091,15 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_create_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_create_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_create_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookRuntimeTemplateRequest.pb( notebook_service.CreateNotebookRuntimeTemplateRequest() ) @@ -12058,6 +12123,7 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_notebook_runtime_template( request, @@ -12069,6 +12135,7 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_runtime_template_rest_bad_request( @@ -12174,10 +12241,14 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_runtime_template" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_runtime_template" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeTemplateRequest.pb( notebook_service.GetNotebookRuntimeTemplateRequest() ) @@ -12203,6 +12274,10 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) client.get_notebook_runtime_template( request, @@ -12214,6 +12289,7 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_runtime_templates_rest_bad_request( @@ -12301,10 +12377,14 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_list_notebook_runtime_templates", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_runtime_templates_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_runtime_templates" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimeTemplatesRequest.pb( notebook_service.ListNotebookRuntimeTemplatesRequest() ) @@ -12330,6 +12410,10 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimeTemplatesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimeTemplatesResponse(), + metadata, + ) client.list_notebook_runtime_templates( request, @@ -12341,6 +12425,7 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_runtime_template_rest_bad_request( @@ -12426,11 +12511,15 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_delete_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeTemplateRequest.pb( notebook_service.DeleteNotebookRuntimeTemplateRequest() ) @@ -12454,6 +12543,7 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_runtime_template( request, @@ -12465,6 +12555,7 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_notebook_runtime_template_rest_bad_request( @@ -12553,6 +12644,14 @@ def test_update_notebook_runtime_template_rest_call_success(request_type): "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -12688,11 +12787,15 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_update_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_update_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_update_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpdateNotebookRuntimeTemplateRequest.pb( notebook_service.UpdateNotebookRuntimeTemplateRequest() ) @@ -12718,6 +12821,10 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) client.update_notebook_runtime_template( request, @@ -12729,6 +12836,7 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_assign_notebook_runtime_rest_bad_request( @@ -12809,10 +12917,14 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_assign_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_assign_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_assign_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.AssignNotebookRuntimeRequest.pb( notebook_service.AssignNotebookRuntimeRequest() ) @@ -12836,6 +12948,7 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.assign_notebook_runtime( request, @@ -12847,6 +12960,7 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_runtime_rest_bad_request( @@ -12966,10 +13080,14 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeRequest.pb( notebook_service.GetNotebookRuntimeRequest() ) @@ -12995,6 +13113,7 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntime() + post_with_metadata.return_value = notebook_runtime.NotebookRuntime(), metadata client.get_notebook_runtime( request, @@ -13006,6 +13125,7 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_runtimes_rest_bad_request( @@ -13090,10 +13210,14 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_list_notebook_runtimes" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_runtimes_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_runtimes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimesRequest.pb( notebook_service.ListNotebookRuntimesRequest() ) @@ -13119,6 +13243,10 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimesResponse(), + metadata, + ) client.list_notebook_runtimes( request, @@ -13130,6 +13258,7 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_runtime_rest_bad_request( @@ -13214,10 +13343,14 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_delete_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeRequest.pb( notebook_service.DeleteNotebookRuntimeRequest() ) @@ -13241,6 +13374,7 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_runtime( request, @@ -13252,6 +13386,7 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upgrade_notebook_runtime_rest_bad_request( @@ -13336,10 +13471,14 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_upgrade_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_upgrade_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_upgrade_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpgradeNotebookRuntimeRequest.pb( notebook_service.UpgradeNotebookRuntimeRequest() ) @@ -13363,6 +13502,7 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.upgrade_notebook_runtime( request, @@ -13374,6 +13514,7 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_start_notebook_runtime_rest_bad_request( @@ -13458,10 +13599,14 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_start_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_start_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_start_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StartNotebookRuntimeRequest.pb( notebook_service.StartNotebookRuntimeRequest() ) @@ -13485,6 +13630,7 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.start_notebook_runtime( request, @@ -13496,6 +13642,7 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stop_notebook_runtime_rest_bad_request( @@ -13580,10 +13727,14 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_stop_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_stop_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_stop_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StopNotebookRuntimeRequest.pb( notebook_service.StopNotebookRuntimeRequest() ) @@ -13607,6 +13758,7 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.stop_notebook_runtime( request, @@ -13618,6 +13770,7 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_notebook_execution_job_rest_bad_request( @@ -13824,10 +13977,14 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_create_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_create_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_create_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookExecutionJobRequest.pb( notebook_service.CreateNotebookExecutionJobRequest() ) @@ -13851,6 +14008,7 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_notebook_execution_job( request, @@ -13862,6 +14020,7 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_execution_job_rest_bad_request( @@ -13961,10 +14120,14 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookExecutionJobRequest.pb( notebook_service.GetNotebookExecutionJobRequest() ) @@ -13990,6 +14153,10 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_execution_job.NotebookExecutionJob() + post_with_metadata.return_value = ( + notebook_execution_job.NotebookExecutionJob(), + metadata, + ) client.get_notebook_execution_job( request, @@ -14001,6 +14168,7 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_execution_jobs_rest_bad_request( @@ -14087,10 +14255,14 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_list_notebook_execution_jobs" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_execution_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_execution_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookExecutionJobsRequest.pb( notebook_service.ListNotebookExecutionJobsRequest() ) @@ -14116,6 +14288,10 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookExecutionJobsResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookExecutionJobsResponse(), + metadata, + ) client.list_notebook_execution_jobs( request, @@ -14127,6 +14303,7 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_execution_job_rest_bad_request( @@ -14211,10 +14388,14 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_delete_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookExecutionJobRequest.pb( notebook_service.DeleteNotebookExecutionJobRequest() ) @@ -14238,6 +14419,7 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_execution_job( request, @@ -14249,6 +14431,7 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -15350,6 +15533,14 @@ async def test_create_notebook_runtime_template_rest_asyncio_call_success(reques "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -15473,11 +15664,15 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_create_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_create_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_create_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookRuntimeTemplateRequest.pb( notebook_service.CreateNotebookRuntimeTemplateRequest() ) @@ -15501,6 +15696,7 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_notebook_runtime_template( request, @@ -15512,6 +15708,7 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15636,11 +15833,15 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeTemplateRequest.pb( notebook_service.GetNotebookRuntimeTemplateRequest() ) @@ -15666,6 +15867,10 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) await client.get_notebook_runtime_template( request, @@ -15677,6 +15882,7 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15782,11 +15988,15 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_runtime_templates", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_runtime_templates_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_runtime_templates", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimeTemplatesRequest.pb( notebook_service.ListNotebookRuntimeTemplatesRequest() ) @@ -15812,6 +16022,10 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimeTemplatesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimeTemplatesResponse(), + metadata, + ) await client.list_notebook_runtime_templates( request, @@ -15823,6 +16037,7 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15926,11 +16141,15 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeTemplateRequest.pb( notebook_service.DeleteNotebookRuntimeTemplateRequest() ) @@ -15954,6 +16173,7 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_runtime_template( request, @@ -15965,6 +16185,7 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16062,6 +16283,14 @@ async def test_update_notebook_runtime_template_rest_asyncio_call_success(reques "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -16206,11 +16435,15 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_update_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_update_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_update_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpdateNotebookRuntimeTemplateRequest.pb( notebook_service.UpdateNotebookRuntimeTemplateRequest() ) @@ -16236,6 +16469,10 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) await client.update_notebook_runtime_template( request, @@ -16247,6 +16484,7 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16343,10 +16581,14 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_assign_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_assign_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_assign_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.AssignNotebookRuntimeRequest.pb( notebook_service.AssignNotebookRuntimeRequest() ) @@ -16370,6 +16612,7 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.assign_notebook_runtime( request, @@ -16381,6 +16624,7 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16516,10 +16760,14 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeRequest.pb( notebook_service.GetNotebookRuntimeRequest() ) @@ -16545,6 +16793,7 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntime() + post_with_metadata.return_value = notebook_runtime.NotebookRuntime(), metadata await client.get_notebook_runtime( request, @@ -16556,6 +16805,7 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16656,10 +16906,14 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_runtimes" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_runtimes_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_runtimes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimesRequest.pb( notebook_service.ListNotebookRuntimesRequest() ) @@ -16685,6 +16939,10 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimesResponse(), + metadata, + ) await client.list_notebook_runtimes( request, @@ -16696,6 +16954,7 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16796,10 +17055,14 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeRequest.pb( notebook_service.DeleteNotebookRuntimeRequest() ) @@ -16823,6 +17086,7 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_runtime( request, @@ -16834,6 +17098,7 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16934,10 +17199,14 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_upgrade_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_upgrade_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_upgrade_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpgradeNotebookRuntimeRequest.pb( notebook_service.UpgradeNotebookRuntimeRequest() ) @@ -16961,6 +17230,7 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.upgrade_notebook_runtime( request, @@ -16972,6 +17242,7 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17072,10 +17343,14 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_start_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_start_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_start_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StartNotebookRuntimeRequest.pb( notebook_service.StartNotebookRuntimeRequest() ) @@ -17099,6 +17374,7 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.start_notebook_runtime( request, @@ -17110,6 +17386,7 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17210,10 +17487,14 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_stop_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_stop_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_stop_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StopNotebookRuntimeRequest.pb( notebook_service.StopNotebookRuntimeRequest() ) @@ -17237,6 +17518,7 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.stop_notebook_runtime( request, @@ -17248,6 +17530,7 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17473,11 +17756,15 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_create_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_create_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_create_notebook_execution_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookExecutionJobRequest.pb( notebook_service.CreateNotebookExecutionJobRequest() ) @@ -17501,6 +17788,7 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_notebook_execution_job( request, @@ -17512,6 +17800,7 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17628,10 +17917,14 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookExecutionJobRequest.pb( notebook_service.GetNotebookExecutionJobRequest() ) @@ -17657,6 +17950,10 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = notebook_execution_job.NotebookExecutionJob() + post_with_metadata.return_value = ( + notebook_execution_job.NotebookExecutionJob(), + metadata, + ) await client.get_notebook_execution_job( request, @@ -17668,6 +17965,7 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17771,11 +18069,15 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_execution_jobs", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_execution_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_execution_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookExecutionJobsRequest.pb( notebook_service.ListNotebookExecutionJobsRequest() ) @@ -17801,6 +18103,10 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookExecutionJobsResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookExecutionJobsResponse(), + metadata, + ) await client.list_notebook_execution_jobs( request, @@ -17812,6 +18118,7 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17915,11 +18222,15 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_execution_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookExecutionJobRequest.pb( notebook_service.DeleteNotebookExecutionJobRequest() ) @@ -17943,6 +18254,7 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_execution_job( request, @@ -17954,6 +18266,7 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_persistent_resource_service.py b/tests/unit/gapic/aiplatform_v1/test_persistent_resource_service.py index 9ddb342010..c3ac8c1c40 100644 --- a/tests/unit/gapic/aiplatform_v1/test_persistent_resource_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_persistent_resource_service.py @@ -98,6 +98,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -376,6 +384,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PersistentResourceServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PersistentResourceServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5415,11 +5466,15 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_create_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_create_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_create_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.CreatePersistentResourceRequest.pb( persistent_resource_service.CreatePersistentResourceRequest() ) @@ -5443,6 +5498,7 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_persistent_resource( request, @@ -5454,6 +5510,7 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_persistent_resource_rest_bad_request( @@ -5551,11 +5608,15 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_get_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_get_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_get_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.GetPersistentResourceRequest.pb( persistent_resource_service.GetPersistentResourceRequest() ) @@ -5581,6 +5642,10 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = persistent_resource.PersistentResource() + post_with_metadata.return_value = ( + persistent_resource.PersistentResource(), + metadata, + ) client.get_persistent_resource( request, @@ -5592,6 +5657,7 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_persistent_resources_rest_bad_request( @@ -5679,11 +5745,15 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_list_persistent_resources", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_list_persistent_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_list_persistent_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.ListPersistentResourcesRequest.pb( persistent_resource_service.ListPersistentResourcesRequest() ) @@ -5713,6 +5783,10 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): post.return_value = ( persistent_resource_service.ListPersistentResourcesResponse() ) + post_with_metadata.return_value = ( + persistent_resource_service.ListPersistentResourcesResponse(), + metadata, + ) client.list_persistent_resources( request, @@ -5724,6 +5798,7 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_persistent_resource_rest_bad_request( @@ -5809,11 +5884,15 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_delete_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_delete_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_delete_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.DeletePersistentResourceRequest.pb( persistent_resource_service.DeletePersistentResourceRequest() ) @@ -5837,6 +5916,7 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_persistent_resource( request, @@ -5848,6 +5928,7 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_persistent_resource_rest_bad_request( @@ -6073,11 +6154,15 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_update_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_update_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_update_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.UpdatePersistentResourceRequest.pb( persistent_resource_service.UpdatePersistentResourceRequest() ) @@ -6101,6 +6186,7 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_persistent_resource( request, @@ -6112,6 +6198,7 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_reboot_persistent_resource_rest_bad_request( @@ -6197,11 +6284,15 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_reboot_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_reboot_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_reboot_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.RebootPersistentResourceRequest.pb( persistent_resource_service.RebootPersistentResourceRequest() ) @@ -6225,6 +6316,7 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.reboot_persistent_resource( request, @@ -6236,6 +6328,7 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7261,11 +7354,15 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_create_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_create_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_create_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.CreatePersistentResourceRequest.pb( persistent_resource_service.CreatePersistentResourceRequest() ) @@ -7289,6 +7386,7 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_persistent_resource( request, @@ -7300,6 +7398,7 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7413,11 +7512,15 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto transports.AsyncPersistentResourceServiceRestInterceptor, "post_get_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_get_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_get_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.GetPersistentResourceRequest.pb( persistent_resource_service.GetPersistentResourceRequest() ) @@ -7443,6 +7546,10 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = persistent_resource.PersistentResource() + post_with_metadata.return_value = ( + persistent_resource.PersistentResource(), + metadata, + ) await client.get_persistent_resource( request, @@ -7454,6 +7561,7 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7557,11 +7665,15 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep transports.AsyncPersistentResourceServiceRestInterceptor, "post_list_persistent_resources", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_list_persistent_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_list_persistent_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.ListPersistentResourcesRequest.pb( persistent_resource_service.ListPersistentResourcesRequest() ) @@ -7591,6 +7703,10 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep post.return_value = ( persistent_resource_service.ListPersistentResourcesResponse() ) + post_with_metadata.return_value = ( + persistent_resource_service.ListPersistentResourcesResponse(), + metadata, + ) await client.list_persistent_resources( request, @@ -7602,6 +7718,7 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7703,11 +7820,15 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_delete_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_delete_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_delete_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.DeletePersistentResourceRequest.pb( persistent_resource_service.DeletePersistentResourceRequest() ) @@ -7731,6 +7852,7 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_persistent_resource( request, @@ -7742,6 +7864,7 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7983,11 +8106,15 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_update_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_update_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_update_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.UpdatePersistentResourceRequest.pb( persistent_resource_service.UpdatePersistentResourceRequest() ) @@ -8011,6 +8138,7 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_persistent_resource( request, @@ -8022,6 +8150,7 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8123,11 +8252,15 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_reboot_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_reboot_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_reboot_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.RebootPersistentResourceRequest.pb( persistent_resource_service.RebootPersistentResourceRequest() ) @@ -8151,6 +8284,7 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.reboot_persistent_resource( request, @@ -8162,6 +8296,7 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_pipeline_service.py b/tests/unit/gapic/aiplatform_v1/test_pipeline_service.py index eca386f2df..27c41d529e 100644 --- a/tests/unit/gapic/aiplatform_v1/test_pipeline_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_pipeline_service.py @@ -107,6 +107,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -364,6 +372,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PipelineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PipelineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -9049,6 +9100,7 @@ def test_create_training_pipeline_rest_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -9074,10 +9126,25 @@ def test_create_training_pipeline_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -9298,10 +9365,14 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_create_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_create_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_create_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreateTrainingPipelineRequest.pb( pipeline_service.CreateTrainingPipelineRequest() ) @@ -9327,6 +9398,10 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_training_pipeline.TrainingPipeline() + post_with_metadata.return_value = ( + gca_training_pipeline.TrainingPipeline(), + metadata, + ) client.create_training_pipeline( request, @@ -9338,6 +9413,7 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_training_pipeline_rest_bad_request( @@ -9436,10 +9512,14 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_get_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_get_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_get_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetTrainingPipelineRequest.pb( pipeline_service.GetTrainingPipelineRequest() ) @@ -9465,6 +9545,7 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = training_pipeline.TrainingPipeline() + post_with_metadata.return_value = training_pipeline.TrainingPipeline(), metadata client.get_training_pipeline( request, @@ -9476,6 +9557,7 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_training_pipelines_rest_bad_request( @@ -9560,10 +9642,14 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_list_training_pipelines" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_list_training_pipelines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_list_training_pipelines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListTrainingPipelinesRequest.pb( pipeline_service.ListTrainingPipelinesRequest() ) @@ -9589,6 +9675,10 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListTrainingPipelinesResponse() + post_with_metadata.return_value = ( + pipeline_service.ListTrainingPipelinesResponse(), + metadata, + ) client.list_training_pipelines( request, @@ -9600,6 +9690,7 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_training_pipeline_rest_bad_request( @@ -9684,10 +9775,14 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_delete_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_delete_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_delete_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeleteTrainingPipelineRequest.pb( pipeline_service.DeleteTrainingPipelineRequest() ) @@ -9711,6 +9806,7 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_training_pipeline( request, @@ -9722,6 +9818,7 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_training_pipeline_rest_bad_request( @@ -10105,10 +10202,14 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_create_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_create_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_create_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreatePipelineJobRequest.pb( pipeline_service.CreatePipelineJobRequest() ) @@ -10134,6 +10235,7 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_pipeline_job.PipelineJob() + post_with_metadata.return_value = gca_pipeline_job.PipelineJob(), metadata client.create_pipeline_job( request, @@ -10145,6 +10247,7 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_pipeline_job_rest_bad_request( @@ -10245,10 +10348,13 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_get_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, "post_get_pipeline_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_get_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetPipelineJobRequest.pb( pipeline_service.GetPipelineJobRequest() ) @@ -10272,6 +10378,7 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_job.PipelineJob() + post_with_metadata.return_value = pipeline_job.PipelineJob(), metadata client.get_pipeline_job( request, @@ -10283,6 +10390,7 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_pipeline_jobs_rest_bad_request( @@ -10367,10 +10475,14 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_list_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_list_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_list_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListPipelineJobsRequest.pb( pipeline_service.ListPipelineJobsRequest() ) @@ -10396,6 +10508,10 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListPipelineJobsResponse() + post_with_metadata.return_value = ( + pipeline_service.ListPipelineJobsResponse(), + metadata, + ) client.list_pipeline_jobs( request, @@ -10407,6 +10523,7 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_pipeline_job_rest_bad_request( @@ -10487,10 +10604,14 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_delete_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_delete_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_delete_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeletePipelineJobRequest.pb( pipeline_service.DeletePipelineJobRequest() ) @@ -10514,6 +10635,7 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_pipeline_job( request, @@ -10525,6 +10647,7 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_delete_pipeline_jobs_rest_bad_request( @@ -10605,10 +10728,14 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_batch_delete_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_batch_delete_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_batch_delete_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchDeletePipelineJobsRequest.pb( pipeline_service.BatchDeletePipelineJobsRequest() ) @@ -10632,6 +10759,7 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_delete_pipeline_jobs( request, @@ -10643,6 +10771,7 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_pipeline_job_rest_bad_request( @@ -10832,10 +10961,14 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_batch_cancel_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_batch_cancel_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_batch_cancel_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchCancelPipelineJobsRequest.pb( pipeline_service.BatchCancelPipelineJobsRequest() ) @@ -10859,6 +10992,7 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_cancel_pipeline_jobs( request, @@ -10870,6 +11004,7 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -11896,6 +12031,7 @@ async def test_create_training_pipeline_rest_asyncio_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -11921,10 +12057,25 @@ async def test_create_training_pipeline_rest_asyncio_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -12152,10 +12303,14 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_create_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_create_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_create_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreateTrainingPipelineRequest.pb( pipeline_service.CreateTrainingPipelineRequest() ) @@ -12181,6 +12336,10 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = gca_training_pipeline.TrainingPipeline() + post_with_metadata.return_value = ( + gca_training_pipeline.TrainingPipeline(), + metadata, + ) await client.create_training_pipeline( request, @@ -12192,6 +12351,7 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12306,10 +12466,14 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_get_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_get_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_get_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetTrainingPipelineRequest.pb( pipeline_service.GetTrainingPipelineRequest() ) @@ -12335,6 +12499,7 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = training_pipeline.TrainingPipeline() + post_with_metadata.return_value = training_pipeline.TrainingPipeline(), metadata await client.get_training_pipeline( request, @@ -12346,6 +12511,7 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12446,10 +12612,14 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_list_training_pipelines" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_list_training_pipelines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_list_training_pipelines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListTrainingPipelinesRequest.pb( pipeline_service.ListTrainingPipelinesRequest() ) @@ -12475,6 +12645,10 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = pipeline_service.ListTrainingPipelinesResponse() + post_with_metadata.return_value = ( + pipeline_service.ListTrainingPipelinesResponse(), + metadata, + ) await client.list_training_pipelines( request, @@ -12486,6 +12660,7 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12586,10 +12761,14 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_delete_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_delete_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_delete_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeleteTrainingPipelineRequest.pb( pipeline_service.DeleteTrainingPipelineRequest() ) @@ -12613,6 +12792,7 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_training_pipeline( request, @@ -12624,6 +12804,7 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13039,10 +13220,14 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_create_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_create_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_create_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreatePipelineJobRequest.pb( pipeline_service.CreatePipelineJobRequest() ) @@ -13068,6 +13253,7 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_pipeline_job.PipelineJob() + post_with_metadata.return_value = gca_pipeline_job.PipelineJob(), metadata await client.create_pipeline_job( request, @@ -13079,6 +13265,7 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13195,10 +13382,14 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_get_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_get_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_get_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetPipelineJobRequest.pb( pipeline_service.GetPipelineJobRequest() ) @@ -13222,6 +13413,7 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_job.PipelineJob() + post_with_metadata.return_value = pipeline_job.PipelineJob(), metadata await client.get_pipeline_job( request, @@ -13233,6 +13425,7 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13333,10 +13526,14 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_list_pipeline_jobs" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_list_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_list_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListPipelineJobsRequest.pb( pipeline_service.ListPipelineJobsRequest() ) @@ -13362,6 +13559,10 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListPipelineJobsResponse() + post_with_metadata.return_value = ( + pipeline_service.ListPipelineJobsResponse(), + metadata, + ) await client.list_pipeline_jobs( request, @@ -13373,6 +13574,7 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13469,10 +13671,14 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_delete_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_delete_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_delete_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeletePipelineJobRequest.pb( pipeline_service.DeletePipelineJobRequest() ) @@ -13496,6 +13702,7 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_pipeline_job( request, @@ -13507,6 +13714,7 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13604,10 +13812,14 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce transports.AsyncPipelineServiceRestInterceptor, "post_batch_delete_pipeline_jobs", ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_batch_delete_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_batch_delete_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchDeletePipelineJobsRequest.pb( pipeline_service.BatchDeletePipelineJobsRequest() ) @@ -13631,6 +13843,7 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_delete_pipeline_jobs( request, @@ -13642,6 +13855,7 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13864,10 +14078,14 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce transports.AsyncPipelineServiceRestInterceptor, "post_batch_cancel_pipeline_jobs", ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_batch_cancel_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_batch_cancel_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchCancelPipelineJobsRequest.pb( pipeline_service.BatchCancelPipelineJobsRequest() ) @@ -13891,6 +14109,7 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_cancel_pipeline_jobs( request, @@ -13902,6 +14121,7 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_prediction_service.py b/tests/unit/gapic/aiplatform_v1/test_prediction_service.py index b95dcc7ced..3a5bd3daed 100644 --- a/tests/unit/gapic/aiplatform_v1/test_prediction_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_prediction_service.py @@ -86,10 +86,19 @@ from google.protobuf import any_pb2 # type: ignore from google.protobuf import duration_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from google.type import latlng_pb2 # type: ignore import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -356,6 +365,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PredictionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PredictionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3775,6 +3827,7 @@ def test_generate_content(request_type, transport: str = "grpc"): # Designate an appropriate return value for the call. call.return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) response = client.generate_content(request) @@ -3787,6 +3840,7 @@ def test_generate_content(request_type, transport: str = "grpc"): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" def test_generate_content_non_empty_request_with_auto_populated_field(): @@ -3918,6 +3972,7 @@ async def test_generate_content_async( call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) ) response = await client.generate_content(request) @@ -3931,6 +3986,7 @@ async def test_generate_content_async( # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -6585,6 +6641,7 @@ async def test_generate_content_empty_call_grpc_asyncio(): call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) ) await client.generate_content(request=None) @@ -6718,10 +6775,13 @@ def test_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.PredictRequest.pb( prediction_service.PredictRequest() ) @@ -6747,6 +6807,7 @@ def test_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.PredictResponse() + post_with_metadata.return_value = prediction_service.PredictResponse(), metadata client.predict( request, @@ -6758,6 +6819,7 @@ def test_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_raw_predict_rest_bad_request( @@ -6841,10 +6903,13 @@ def test_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_raw_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.RawPredictRequest.pb( prediction_service.RawPredictRequest() ) @@ -6868,6 +6933,7 @@ def test_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.raw_predict( request, @@ -6879,6 +6945,7 @@ def test_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_raw_predict_rest_bad_request( @@ -6966,10 +7033,14 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_stream_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_stream_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_stream_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamRawPredictRequest.pb( prediction_service.StreamRawPredictRequest() ) @@ -6993,6 +7064,7 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.stream_raw_predict( request, @@ -7004,6 +7076,7 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_direct_predict_rest_bad_request( @@ -7085,10 +7158,13 @@ def test_direct_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_direct_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_direct_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_direct_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectPredictRequest.pb( prediction_service.DirectPredictRequest() ) @@ -7114,6 +7190,10 @@ def test_direct_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectPredictResponse(), + metadata, + ) client.direct_predict( request, @@ -7125,6 +7205,7 @@ def test_direct_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_direct_raw_predict_rest_bad_request( @@ -7209,10 +7290,14 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_direct_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_direct_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_direct_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectRawPredictRequest.pb( prediction_service.DirectRawPredictRequest() ) @@ -7238,6 +7323,10 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectRawPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectRawPredictResponse(), + metadata, + ) client.direct_raw_predict( request, @@ -7249,6 +7338,7 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_direct_predict_rest_error(): @@ -7373,10 +7463,14 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_server_streaming_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_server_streaming_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_server_streaming_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamingPredictRequest.pb( prediction_service.StreamingPredictRequest() ) @@ -7402,6 +7496,10 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.StreamingPredictResponse() + post_with_metadata.return_value = ( + prediction_service.StreamingPredictResponse(), + metadata, + ) client.server_streaming_predict( request, @@ -7413,6 +7511,7 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_streaming_raw_predict_rest_error(): @@ -7508,10 +7607,13 @@ def test_explain_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_explain" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_explain_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_explain" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ExplainRequest.pb( prediction_service.ExplainRequest() ) @@ -7537,6 +7639,7 @@ def test_explain_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.ExplainResponse() + post_with_metadata.return_value = prediction_service.ExplainResponse(), metadata client.explain( request, @@ -7548,6 +7651,7 @@ def test_explain_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_generate_content_rest_bad_request( @@ -7596,6 +7700,7 @@ def test_generate_content_rest_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -7613,6 +7718,7 @@ def test_generate_content_rest_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.parametrize("null_interceptor", [True, False]) @@ -7632,10 +7738,14 @@ def test_generate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_generate_content" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -7661,6 +7771,10 @@ def test_generate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) client.generate_content( request, @@ -7672,6 +7786,7 @@ def test_generate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_generate_content_rest_bad_request( @@ -7720,6 +7835,7 @@ def test_stream_generate_content_rest_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -7741,6 +7857,7 @@ def test_stream_generate_content_rest_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.parametrize("null_interceptor", [True, False]) @@ -7760,10 +7877,14 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_stream_generate_content" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_stream_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_stream_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -7789,6 +7910,10 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) client.stream_generate_content( request, @@ -7800,6 +7925,7 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -8737,10 +8863,13 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, "post_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.PredictRequest.pb( prediction_service.PredictRequest() ) @@ -8766,6 +8895,7 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.PredictResponse() + post_with_metadata.return_value = prediction_service.PredictResponse(), metadata await client.predict( request, @@ -8777,6 +8907,7 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8876,10 +9007,14 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.RawPredictRequest.pb( prediction_service.RawPredictRequest() ) @@ -8903,6 +9038,7 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.raw_predict( request, @@ -8914,6 +9050,7 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9015,10 +9152,14 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_stream_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_stream_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_stream_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamRawPredictRequest.pb( prediction_service.StreamRawPredictRequest() ) @@ -9042,6 +9183,7 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.stream_raw_predict( request, @@ -9053,6 +9195,7 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9150,10 +9293,14 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_direct_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_direct_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_direct_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectPredictRequest.pb( prediction_service.DirectPredictRequest() ) @@ -9179,6 +9326,10 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectPredictResponse(), + metadata, + ) await client.direct_predict( request, @@ -9190,6 +9341,7 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9290,10 +9442,14 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_direct_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_direct_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_direct_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectRawPredictRequest.pb( prediction_service.DirectRawPredictRequest() ) @@ -9319,6 +9475,10 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectRawPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectRawPredictResponse(), + metadata, + ) await client.direct_raw_predict( request, @@ -9330,6 +9490,7 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9484,10 +9645,14 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept transports.AsyncPredictionServiceRestInterceptor, "post_server_streaming_predict", ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_server_streaming_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_server_streaming_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamingPredictRequest.pb( prediction_service.StreamingPredictRequest() ) @@ -9513,6 +9678,10 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = prediction_service.StreamingPredictResponse() + post_with_metadata.return_value = ( + prediction_service.StreamingPredictResponse(), + metadata, + ) await client.server_streaming_predict( request, @@ -9524,6 +9693,7 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9642,10 +9812,13 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_explain" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, "post_explain_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_explain" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ExplainRequest.pb( prediction_service.ExplainRequest() ) @@ -9671,6 +9844,7 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.ExplainResponse() + post_with_metadata.return_value = prediction_service.ExplainResponse(), metadata await client.explain( request, @@ -9682,6 +9856,7 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9739,6 +9914,7 @@ async def test_generate_content_rest_asyncio_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -9758,6 +9934,7 @@ async def test_generate_content_rest_asyncio_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -9782,10 +9959,14 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_generate_content" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -9811,6 +9992,10 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) await client.generate_content( request, @@ -9822,6 +10007,7 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9879,6 +10065,7 @@ async def test_stream_generate_content_rest_asyncio_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -9900,6 +10087,7 @@ async def test_stream_generate_content_rest_asyncio_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -9924,10 +10112,14 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_stream_generate_content" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_stream_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_stream_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -9953,6 +10145,10 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) await client.stream_generate_content( request, @@ -9964,6 +10160,7 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_execution_service.py b/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_execution_service.py index ee73172135..579660a87b 100644 --- a/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_execution_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_execution_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -377,6 +385,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ReasoningEngineExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ReasoningEngineExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2278,11 +2329,15 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): transports.ReasoningEngineExecutionServiceRestInterceptor, "post_query_reasoning_engine", ) as post, mock.patch.object( + transports.ReasoningEngineExecutionServiceRestInterceptor, + "post_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineExecutionServiceRestInterceptor, "pre_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_execution_service.QueryReasoningEngineRequest.pb( reasoning_engine_execution_service.QueryReasoningEngineRequest() ) @@ -2312,6 +2367,10 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): post.return_value = ( reasoning_engine_execution_service.QueryReasoningEngineResponse() ) + post_with_metadata.return_value = ( + reasoning_engine_execution_service.QueryReasoningEngineResponse(), + metadata, + ) client.query_reasoning_engine( request, @@ -2323,6 +2382,7 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_query_reasoning_engine_rest_bad_request( @@ -2415,11 +2475,15 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): transports.ReasoningEngineExecutionServiceRestInterceptor, "post_stream_query_reasoning_engine", ) as post, mock.patch.object( + transports.ReasoningEngineExecutionServiceRestInterceptor, + "post_stream_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineExecutionServiceRestInterceptor, "pre_stream_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest.pb( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest() @@ -2445,6 +2509,7 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.stream_query_reasoning_engine( request, @@ -2456,6 +2521,7 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3253,11 +3319,15 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "post_query_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineExecutionServiceRestInterceptor, + "post_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "pre_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_execution_service.QueryReasoningEngineRequest.pb( reasoning_engine_execution_service.QueryReasoningEngineRequest() ) @@ -3287,6 +3357,10 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor post.return_value = ( reasoning_engine_execution_service.QueryReasoningEngineResponse() ) + post_with_metadata.return_value = ( + reasoning_engine_execution_service.QueryReasoningEngineResponse(), + metadata, + ) await client.query_reasoning_engine( request, @@ -3298,6 +3372,7 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3406,11 +3481,15 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "post_stream_query_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineExecutionServiceRestInterceptor, + "post_stream_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "pre_stream_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest.pb( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest() @@ -3436,6 +3515,7 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.stream_query_reasoning_engine( request, @@ -3447,6 +3527,7 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_service.py b/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_service.py index a661d7c9e4..8be2ad59ae 100644 --- a/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_reasoning_engine_service.py @@ -91,6 +91,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -365,6 +373,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ReasoningEngineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ReasoningEngineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4677,10 +4728,14 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_create_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_create_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_create_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.CreateReasoningEngineRequest.pb( reasoning_engine_service.CreateReasoningEngineRequest() ) @@ -4704,6 +4759,7 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_reasoning_engine( request, @@ -4715,6 +4771,7 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_reasoning_engine_rest_bad_request( @@ -4809,10 +4866,14 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_get_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_get_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_get_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.GetReasoningEngineRequest.pb( reasoning_engine_service.GetReasoningEngineRequest() ) @@ -4838,6 +4899,7 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine.ReasoningEngine() + post_with_metadata.return_value = reasoning_engine.ReasoningEngine(), metadata client.get_reasoning_engine( request, @@ -4849,6 +4911,7 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_reasoning_engines_rest_bad_request( @@ -4935,10 +4998,14 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_list_reasoning_engines" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_list_reasoning_engines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_list_reasoning_engines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.ListReasoningEnginesRequest.pb( reasoning_engine_service.ListReasoningEnginesRequest() ) @@ -4964,6 +5031,10 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine_service.ListReasoningEnginesResponse() + post_with_metadata.return_value = ( + reasoning_engine_service.ListReasoningEnginesResponse(), + metadata, + ) client.list_reasoning_engines( request, @@ -4975,6 +5046,7 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_reasoning_engine_rest_bad_request( @@ -5149,10 +5221,14 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_update_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_update_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_update_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.UpdateReasoningEngineRequest.pb( reasoning_engine_service.UpdateReasoningEngineRequest() ) @@ -5176,6 +5252,7 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_reasoning_engine( request, @@ -5187,6 +5264,7 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_reasoning_engine_rest_bad_request( @@ -5271,10 +5349,14 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_delete_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_delete_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_delete_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.DeleteReasoningEngineRequest.pb( reasoning_engine_service.DeleteReasoningEngineRequest() ) @@ -5298,6 +5380,7 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_reasoning_engine( request, @@ -5309,6 +5392,7 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6262,11 +6346,15 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_create_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_create_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_create_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.CreateReasoningEngineRequest.pb( reasoning_engine_service.CreateReasoningEngineRequest() ) @@ -6290,6 +6378,7 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_reasoning_engine( request, @@ -6301,6 +6390,7 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6412,11 +6502,15 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): transports.AsyncReasoningEngineServiceRestInterceptor, "post_get_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_get_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_get_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.GetReasoningEngineRequest.pb( reasoning_engine_service.GetReasoningEngineRequest() ) @@ -6442,6 +6536,7 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine.ReasoningEngine() + post_with_metadata.return_value = reasoning_engine.ReasoningEngine(), metadata await client.get_reasoning_engine( request, @@ -6453,6 +6548,7 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6556,11 +6652,15 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor transports.AsyncReasoningEngineServiceRestInterceptor, "post_list_reasoning_engines", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_list_reasoning_engines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_list_reasoning_engines", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.ListReasoningEnginesRequest.pb( reasoning_engine_service.ListReasoningEnginesRequest() ) @@ -6586,6 +6686,10 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = reasoning_engine_service.ListReasoningEnginesResponse() + post_with_metadata.return_value = ( + reasoning_engine_service.ListReasoningEnginesResponse(), + metadata, + ) await client.list_reasoning_engines( request, @@ -6597,6 +6701,7 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6788,11 +6893,15 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_update_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_update_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_update_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.UpdateReasoningEngineRequest.pb( reasoning_engine_service.UpdateReasoningEngineRequest() ) @@ -6816,6 +6925,7 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_reasoning_engine( request, @@ -6827,6 +6937,7 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6928,11 +7039,15 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_delete_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_delete_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_delete_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.DeleteReasoningEngineRequest.pb( reasoning_engine_service.DeleteReasoningEngineRequest() ) @@ -6956,6 +7071,7 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_reasoning_engine( request, @@ -6967,6 +7083,7 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_schedule_service.py b/tests/unit/gapic/aiplatform_v1/test_schedule_service.py index 9929732b08..1c776e1ae1 100644 --- a/tests/unit/gapic/aiplatform_v1/test_schedule_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_schedule_service.py @@ -108,6 +108,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -365,6 +373,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ScheduleServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ScheduleServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5853,10 +5904,13 @@ def test_create_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_create_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_create_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_create_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.CreateScheduleRequest.pb( schedule_service.CreateScheduleRequest() ) @@ -5880,6 +5934,7 @@ def test_create_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata client.create_schedule( request, @@ -5891,6 +5946,7 @@ def test_create_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_schedule_rest_bad_request( @@ -5971,10 +6027,13 @@ def test_delete_schedule_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_delete_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_delete_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_delete_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.DeleteScheduleRequest.pb( schedule_service.DeleteScheduleRequest() ) @@ -5998,6 +6057,7 @@ def test_delete_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_schedule( request, @@ -6009,6 +6069,7 @@ def test_delete_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_schedule_rest_bad_request( @@ -6108,10 +6169,13 @@ def test_get_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_get_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_get_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_get_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.GetScheduleRequest.pb( schedule_service.GetScheduleRequest() ) @@ -6135,6 +6199,7 @@ def test_get_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule.Schedule() + post_with_metadata.return_value = schedule.Schedule(), metadata client.get_schedule( request, @@ -6146,6 +6211,7 @@ def test_get_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_schedules_rest_bad_request( @@ -6230,10 +6296,13 @@ def test_list_schedules_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_list_schedules" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_list_schedules_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_list_schedules" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.ListSchedulesRequest.pb( schedule_service.ListSchedulesRequest() ) @@ -6259,6 +6328,10 @@ def test_list_schedules_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule_service.ListSchedulesResponse() + post_with_metadata.return_value = ( + schedule_service.ListSchedulesResponse(), + metadata, + ) client.list_schedules( request, @@ -6270,6 +6343,7 @@ def test_list_schedules_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_pause_schedule_rest_bad_request( @@ -6846,10 +6920,13 @@ def test_update_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_update_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_update_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_update_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.UpdateScheduleRequest.pb( schedule_service.UpdateScheduleRequest() ) @@ -6873,6 +6950,7 @@ def test_update_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata client.update_schedule( request, @@ -6884,6 +6962,7 @@ def test_update_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -8054,10 +8133,14 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_create_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_create_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_create_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.CreateScheduleRequest.pb( schedule_service.CreateScheduleRequest() ) @@ -8081,6 +8164,7 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata await client.create_schedule( request, @@ -8092,6 +8176,7 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8188,10 +8273,14 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_delete_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_delete_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_delete_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.DeleteScheduleRequest.pb( schedule_service.DeleteScheduleRequest() ) @@ -8215,6 +8304,7 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_schedule( request, @@ -8226,6 +8316,7 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8341,10 +8432,14 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_get_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_get_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_get_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.GetScheduleRequest.pb( schedule_service.GetScheduleRequest() ) @@ -8368,6 +8463,7 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule.Schedule() + post_with_metadata.return_value = schedule.Schedule(), metadata await client.get_schedule( request, @@ -8379,6 +8475,7 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8479,10 +8576,14 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_list_schedules" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_list_schedules_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_list_schedules" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.ListSchedulesRequest.pb( schedule_service.ListSchedulesRequest() ) @@ -8508,6 +8609,10 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule_service.ListSchedulesResponse() + post_with_metadata.return_value = ( + schedule_service.ListSchedulesResponse(), + metadata, + ) await client.list_schedules( request, @@ -8519,6 +8624,7 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9143,10 +9249,14 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_update_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_update_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_update_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.UpdateScheduleRequest.pb( schedule_service.UpdateScheduleRequest() ) @@ -9170,6 +9280,7 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata await client.update_schedule( request, @@ -9181,6 +9292,7 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_specialist_pool_service.py b/tests/unit/gapic/aiplatform_v1/test_specialist_pool_service.py index fe23695ddc..1e65b14053 100644 --- a/tests/unit/gapic/aiplatform_v1/test_specialist_pool_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_specialist_pool_service.py @@ -89,6 +89,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -363,6 +371,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = SpecialistPoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = SpecialistPoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4691,10 +4742,14 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_create_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_create_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_create_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.CreateSpecialistPoolRequest.pb( specialist_pool_service.CreateSpecialistPoolRequest() ) @@ -4718,6 +4773,7 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_specialist_pool( request, @@ -4729,6 +4785,7 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_specialist_pool_rest_bad_request( @@ -4827,10 +4884,14 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_get_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_get_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_get_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.GetSpecialistPoolRequest.pb( specialist_pool_service.GetSpecialistPoolRequest() ) @@ -4856,6 +4917,7 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool.SpecialistPool() + post_with_metadata.return_value = specialist_pool.SpecialistPool(), metadata client.get_specialist_pool( request, @@ -4867,6 +4929,7 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_specialist_pools_rest_bad_request( @@ -4953,10 +5016,14 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_list_specialist_pools" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_list_specialist_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_list_specialist_pools" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.ListSpecialistPoolsRequest.pb( specialist_pool_service.ListSpecialistPoolsRequest() ) @@ -4982,6 +5049,10 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool_service.ListSpecialistPoolsResponse() + post_with_metadata.return_value = ( + specialist_pool_service.ListSpecialistPoolsResponse(), + metadata, + ) client.list_specialist_pools( request, @@ -4993,6 +5064,7 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_specialist_pool_rest_bad_request( @@ -5077,10 +5149,14 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_delete_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_delete_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_delete_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.DeleteSpecialistPoolRequest.pb( specialist_pool_service.DeleteSpecialistPoolRequest() ) @@ -5104,6 +5180,7 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_specialist_pool( request, @@ -5115,6 +5192,7 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_specialist_pool_rest_bad_request( @@ -5289,10 +5367,14 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_update_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_update_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_update_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.UpdateSpecialistPoolRequest.pb( specialist_pool_service.UpdateSpecialistPoolRequest() ) @@ -5316,6 +5398,7 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_specialist_pool( request, @@ -5327,6 +5410,7 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6280,11 +6364,15 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_create_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_create_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_create_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.CreateSpecialistPoolRequest.pb( specialist_pool_service.CreateSpecialistPoolRequest() ) @@ -6308,6 +6396,7 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_specialist_pool( request, @@ -6319,6 +6408,7 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6433,10 +6523,14 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "post_get_specialist_pool" ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_get_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_get_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.GetSpecialistPoolRequest.pb( specialist_pool_service.GetSpecialistPoolRequest() ) @@ -6462,6 +6556,7 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool.SpecialistPool() + post_with_metadata.return_value = specialist_pool.SpecialistPool(), metadata await client.get_specialist_pool( request, @@ -6473,6 +6568,7 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6576,11 +6672,15 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) transports.AsyncSpecialistPoolServiceRestInterceptor, "post_list_specialist_pools", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_list_specialist_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_list_specialist_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.ListSpecialistPoolsRequest.pb( specialist_pool_service.ListSpecialistPoolsRequest() ) @@ -6606,6 +6706,10 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = specialist_pool_service.ListSpecialistPoolsResponse() + post_with_metadata.return_value = ( + specialist_pool_service.ListSpecialistPoolsResponse(), + metadata, + ) await client.list_specialist_pools( request, @@ -6617,6 +6721,7 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6718,11 +6823,15 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_delete_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_delete_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_delete_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.DeleteSpecialistPoolRequest.pb( specialist_pool_service.DeleteSpecialistPoolRequest() ) @@ -6746,6 +6855,7 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_specialist_pool( request, @@ -6757,6 +6867,7 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6948,11 +7059,15 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_update_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_update_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_update_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.UpdateSpecialistPoolRequest.pb( specialist_pool_service.UpdateSpecialistPoolRequest() ) @@ -6976,6 +7091,7 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_specialist_pool( request, @@ -6987,6 +7103,7 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_tensorboard_service.py b/tests/unit/gapic/aiplatform_v1/test_tensorboard_service.py index a8bcdf8990..b51b0300c1 100644 --- a/tests/unit/gapic/aiplatform_v1/test_tensorboard_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_tensorboard_service.py @@ -102,6 +102,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -372,6 +380,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = TensorboardServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = TensorboardServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -21135,10 +21186,14 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_create_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRequest.pb( tensorboard_service.CreateTensorboardRequest() ) @@ -21162,6 +21217,7 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_tensorboard( request, @@ -21173,6 +21229,7 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_rest_bad_request( @@ -21273,10 +21330,14 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRequest.pb( tensorboard_service.GetTensorboardRequest() ) @@ -21300,6 +21361,7 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard.Tensorboard() + post_with_metadata.return_value = tensorboard.Tensorboard(), metadata client.get_tensorboard( request, @@ -21311,6 +21373,7 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_rest_bad_request( @@ -21481,10 +21544,14 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_update_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRequest.pb( tensorboard_service.UpdateTensorboardRequest() ) @@ -21508,6 +21575,7 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_tensorboard( request, @@ -21519,6 +21587,7 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboards_rest_bad_request( @@ -21603,10 +21672,14 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_list_tensorboards" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboards_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboards" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardsRequest.pb( tensorboard_service.ListTensorboardsRequest() ) @@ -21632,6 +21705,10 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardsResponse(), + metadata, + ) client.list_tensorboards( request, @@ -21643,6 +21720,7 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_rest_bad_request( @@ -21723,10 +21801,14 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRequest.pb( tensorboard_service.DeleteTensorboardRequest() ) @@ -21750,6 +21832,7 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard( request, @@ -21761,6 +21844,7 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_usage_rest_bad_request( @@ -21846,10 +21930,14 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_usage" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_usage_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_usage" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardUsageRequest.pb( tensorboard_service.ReadTensorboardUsageRequest() ) @@ -21875,6 +21963,10 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardUsageResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardUsageResponse(), + metadata, + ) client.read_tensorboard_usage( request, @@ -21886,6 +21978,7 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_size_rest_bad_request( @@ -21974,10 +22067,14 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_size" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_size_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_size" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardSizeRequest.pb( tensorboard_service.ReadTensorboardSizeRequest() ) @@ -22003,6 +22100,10 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardSizeResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardSizeResponse(), + metadata, + ) client.read_tensorboard_size( request, @@ -22014,6 +22115,7 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_experiment_rest_bad_request( @@ -22188,11 +22290,15 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardExperimentRequest.pb( tensorboard_service.CreateTensorboardExperimentRequest() ) @@ -22218,6 +22324,10 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.create_tensorboard_experiment( request, @@ -22229,6 +22339,7 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_experiment_rest_bad_request( @@ -22325,10 +22436,14 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_experiment" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_experiment" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardExperimentRequest.pb( tensorboard_service.GetTensorboardExperimentRequest() ) @@ -22354,6 +22469,10 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.get_tensorboard_experiment( request, @@ -22365,6 +22484,7 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_experiment_rest_bad_request( @@ -22547,11 +22667,15 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardExperimentRequest.pb( tensorboard_service.UpdateTensorboardExperimentRequest() ) @@ -22577,6 +22701,10 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.update_tensorboard_experiment( request, @@ -22588,6 +22716,7 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_experiments_rest_bad_request( @@ -22675,10 +22804,14 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_experiments", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_experiments_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_experiments" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardExperimentsRequest.pb( tensorboard_service.ListTensorboardExperimentsRequest() ) @@ -22704,6 +22837,10 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardExperimentsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardExperimentsResponse(), + metadata, + ) client.list_tensorboard_experiments( request, @@ -22715,6 +22852,7 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_experiment_rest_bad_request( @@ -22800,11 +22938,15 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardExperimentRequest.pb( tensorboard_service.DeleteTensorboardExperimentRequest() ) @@ -22828,6 +22970,7 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_experiment( request, @@ -22839,6 +22982,7 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_run_rest_bad_request( @@ -23011,10 +23155,14 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRunRequest.pb( tensorboard_service.CreateTensorboardRunRequest() ) @@ -23040,6 +23188,7 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata client.create_tensorboard_run( request, @@ -23051,6 +23200,7 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_tensorboard_runs_rest_bad_request( @@ -23139,11 +23289,15 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_batch_create_tensorboard_runs", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_runs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardRunsRequest.pb( tensorboard_service.BatchCreateTensorboardRunsRequest() ) @@ -23169,6 +23323,10 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.BatchCreateTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardRunsResponse(), + metadata, + ) client.batch_create_tensorboard_runs( request, @@ -23180,6 +23338,7 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_run_rest_bad_request( @@ -23274,10 +23433,14 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRunRequest.pb( tensorboard_service.GetTensorboardRunRequest() ) @@ -23303,6 +23466,7 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_run.TensorboardRun() + post_with_metadata.return_value = tensorboard_run.TensorboardRun(), metadata client.get_tensorboard_run( request, @@ -23314,6 +23478,7 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_run_rest_bad_request( @@ -23490,10 +23655,14 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRunRequest.pb( tensorboard_service.UpdateTensorboardRunRequest() ) @@ -23519,6 +23688,7 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata client.update_tensorboard_run( request, @@ -23530,6 +23700,7 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_runs_rest_bad_request( @@ -23618,10 +23789,14 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_runs" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_runs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardRunsRequest.pb( tensorboard_service.ListTensorboardRunsRequest() ) @@ -23647,6 +23822,10 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardRunsResponse(), + metadata, + ) client.list_tensorboard_runs( request, @@ -23658,6 +23837,7 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_run_rest_bad_request( @@ -23742,10 +23922,14 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRunRequest.pb( tensorboard_service.DeleteTensorboardRunRequest() ) @@ -23769,6 +23953,7 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_run( request, @@ -23780,6 +23965,7 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_tensorboard_time_series_rest_bad_request( @@ -23870,11 +24056,15 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor transports.TensorboardServiceRestInterceptor, "post_batch_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardTimeSeriesRequest.pb( tensorboard_service.BatchCreateTensorboardTimeSeriesRequest() ) @@ -23904,6 +24094,10 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor post.return_value = ( tensorboard_service.BatchCreateTensorboardTimeSeriesResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse(), + metadata, + ) client.batch_create_tensorboard_time_series( request, @@ -23915,6 +24109,7 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_time_series_rest_bad_request( @@ -24108,11 +24303,15 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardTimeSeriesRequest.pb( tensorboard_service.CreateTensorboardTimeSeriesRequest() ) @@ -24138,6 +24337,10 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.create_tensorboard_time_series( request, @@ -24149,6 +24352,7 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_time_series_rest_bad_request( @@ -24252,10 +24456,14 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_time_series" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_time_series" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardTimeSeriesRequest.pb( tensorboard_service.GetTensorboardTimeSeriesRequest() ) @@ -24281,6 +24489,10 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.get_tensorboard_time_series( request, @@ -24292,6 +24504,7 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_time_series_rest_bad_request( @@ -24489,11 +24702,15 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardTimeSeriesRequest.pb( tensorboard_service.UpdateTensorboardTimeSeriesRequest() ) @@ -24519,6 +24736,10 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.update_tensorboard_time_series( request, @@ -24530,6 +24751,7 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_time_series_rest_bad_request( @@ -24621,10 +24843,14 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_time_series" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardTimeSeriesRequest.pb( tensorboard_service.ListTensorboardTimeSeriesRequest() ) @@ -24650,6 +24876,10 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardTimeSeriesResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardTimeSeriesResponse(), + metadata, + ) client.list_tensorboard_time_series( request, @@ -24661,6 +24891,7 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_time_series_rest_bad_request( @@ -24746,11 +24977,15 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardTimeSeriesRequest.pb( tensorboard_service.DeleteTensorboardTimeSeriesRequest() ) @@ -24774,6 +25009,7 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_time_series( request, @@ -24785,6 +25021,7 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_read_tensorboard_time_series_data_rest_bad_request( @@ -24877,11 +25114,15 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep transports.TensorboardServiceRestInterceptor, "post_batch_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest() ) @@ -24911,6 +25152,10 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep post.return_value = ( tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse(), + metadata, + ) client.batch_read_tensorboard_time_series_data( request, @@ -24922,6 +25167,7 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_time_series_data_rest_bad_request( @@ -25012,11 +25258,15 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ReadTensorboardTimeSeriesDataRequest() ) @@ -25044,6 +25294,10 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardTimeSeriesDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardTimeSeriesDataResponse(), + metadata, + ) client.read_tensorboard_time_series_data( request, @@ -25055,6 +25309,7 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_blob_data_rest_bad_request( @@ -25146,10 +25401,14 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_blob_data" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_blob_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_blob_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardBlobDataRequest.pb( tensorboard_service.ReadTensorboardBlobDataRequest() ) @@ -25175,6 +25434,10 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardBlobDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardBlobDataResponse(), + metadata, + ) client.read_tensorboard_blob_data( request, @@ -25186,6 +25449,7 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_tensorboard_experiment_data_rest_bad_request( @@ -25276,11 +25540,15 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_write_tensorboard_experiment_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_write_tensorboard_experiment_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_write_tensorboard_experiment_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardExperimentDataRequest.pb( tensorboard_service.WriteTensorboardExperimentDataRequest() ) @@ -25308,6 +25576,10 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardExperimentDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardExperimentDataResponse(), + metadata, + ) client.write_tensorboard_experiment_data( request, @@ -25319,6 +25591,7 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_tensorboard_run_data_rest_bad_request( @@ -25406,10 +25679,14 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_write_tensorboard_run_data" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_write_tensorboard_run_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_write_tensorboard_run_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardRunDataRequest.pb( tensorboard_service.WriteTensorboardRunDataRequest() ) @@ -25435,6 +25712,10 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardRunDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardRunDataResponse(), + metadata, + ) client.write_tensorboard_run_data( request, @@ -25446,6 +25727,7 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_tensorboard_time_series_data_rest_bad_request( @@ -25537,11 +25819,15 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) transports.TensorboardServiceRestInterceptor, "post_export_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_export_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_export_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ExportTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ExportTensorboardTimeSeriesDataRequest() ) @@ -25571,6 +25857,10 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) post.return_value = ( tensorboard_service.ExportTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.ExportTensorboardTimeSeriesDataResponse(), + metadata, + ) client.export_tensorboard_time_series_data( request, @@ -25582,6 +25872,7 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -27078,10 +27369,14 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRequest.pb( tensorboard_service.CreateTensorboardRequest() ) @@ -27105,6 +27400,7 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_tensorboard( request, @@ -27116,6 +27412,7 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27232,10 +27529,14 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRequest.pb( tensorboard_service.GetTensorboardRequest() ) @@ -27259,6 +27560,7 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard.Tensorboard() + post_with_metadata.return_value = tensorboard.Tensorboard(), metadata await client.get_tensorboard( request, @@ -27270,6 +27572,7 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27456,10 +27759,14 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRequest.pb( tensorboard_service.UpdateTensorboardRequest() ) @@ -27483,6 +27790,7 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_tensorboard( request, @@ -27494,6 +27802,7 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27594,10 +27903,14 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboards" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboards_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboards" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardsRequest.pb( tensorboard_service.ListTensorboardsRequest() ) @@ -27623,6 +27936,10 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardsResponse(), + metadata, + ) await client.list_tensorboards( request, @@ -27634,6 +27951,7 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27730,10 +28048,14 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRequest.pb( tensorboard_service.DeleteTensorboardRequest() ) @@ -27757,6 +28079,7 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard( request, @@ -27768,6 +28091,7 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27869,10 +28193,14 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_usage" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_usage_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_usage" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardUsageRequest.pb( tensorboard_service.ReadTensorboardUsageRequest() ) @@ -27898,6 +28226,10 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardUsageResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardUsageResponse(), + metadata, + ) await client.read_tensorboard_usage( request, @@ -27909,6 +28241,7 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28013,10 +28346,14 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_size" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_size_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_size" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardSizeRequest.pb( tensorboard_service.ReadTensorboardSizeRequest() ) @@ -28042,6 +28379,10 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardSizeResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardSizeResponse(), + metadata, + ) await client.read_tensorboard_size( request, @@ -28053,6 +28394,7 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28245,11 +28587,15 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardExperimentRequest.pb( tensorboard_service.CreateTensorboardExperimentRequest() ) @@ -28275,6 +28621,10 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.create_tensorboard_experiment( request, @@ -28286,6 +28636,7 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28399,11 +28750,15 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardExperimentRequest.pb( tensorboard_service.GetTensorboardExperimentRequest() ) @@ -28429,6 +28784,10 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.get_tensorboard_experiment( request, @@ -28440,6 +28799,7 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28640,11 +29000,15 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardExperimentRequest.pb( tensorboard_service.UpdateTensorboardExperimentRequest() ) @@ -28670,6 +29034,10 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.update_tensorboard_experiment( request, @@ -28681,6 +29049,7 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28784,11 +29153,15 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_experiments", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_experiments_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_experiments", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardExperimentsRequest.pb( tensorboard_service.ListTensorboardExperimentsRequest() ) @@ -28814,6 +29187,10 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardExperimentsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardExperimentsResponse(), + metadata, + ) await client.list_tensorboard_experiments( request, @@ -28825,6 +29202,7 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28928,11 +29306,15 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardExperimentRequest.pb( tensorboard_service.DeleteTensorboardExperimentRequest() ) @@ -28956,6 +29338,7 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_experiment( request, @@ -28967,6 +29350,7 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29155,10 +29539,14 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRunRequest.pb( tensorboard_service.CreateTensorboardRunRequest() ) @@ -29184,6 +29572,7 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata await client.create_tensorboard_run( request, @@ -29195,6 +29584,7 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29301,11 +29691,15 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_batch_create_tensorboard_runs", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_runs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardRunsRequest.pb( tensorboard_service.BatchCreateTensorboardRunsRequest() ) @@ -29331,6 +29725,10 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.BatchCreateTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardRunsResponse(), + metadata, + ) await client.batch_create_tensorboard_runs( request, @@ -29342,6 +29740,7 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29452,10 +29851,14 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRunRequest.pb( tensorboard_service.GetTensorboardRunRequest() ) @@ -29481,6 +29884,7 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_run.TensorboardRun() + post_with_metadata.return_value = tensorboard_run.TensorboardRun(), metadata await client.get_tensorboard_run( request, @@ -29492,6 +29896,7 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29684,10 +30089,14 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRunRequest.pb( tensorboard_service.UpdateTensorboardRunRequest() ) @@ -29713,6 +30122,7 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata await client.update_tensorboard_run( request, @@ -29724,6 +30134,7 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29828,10 +30239,14 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_runs" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_runs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardRunsRequest.pb( tensorboard_service.ListTensorboardRunsRequest() ) @@ -29857,6 +30272,10 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardRunsResponse(), + metadata, + ) await client.list_tensorboard_runs( request, @@ -29868,6 +30287,7 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29968,10 +30388,14 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRunRequest.pb( tensorboard_service.DeleteTensorboardRunRequest() ) @@ -29995,6 +30419,7 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_run( request, @@ -30006,6 +30431,7 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30116,11 +30542,15 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_batch_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardTimeSeriesRequest.pb( tensorboard_service.BatchCreateTensorboardTimeSeriesRequest() ) @@ -30150,6 +30580,10 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( post.return_value = ( tensorboard_service.BatchCreateTensorboardTimeSeriesResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse(), + metadata, + ) await client.batch_create_tensorboard_time_series( request, @@ -30161,6 +30595,7 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30372,11 +30807,15 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardTimeSeriesRequest.pb( tensorboard_service.CreateTensorboardTimeSeriesRequest() ) @@ -30402,6 +30841,10 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.create_tensorboard_time_series( request, @@ -30413,6 +30856,7 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30533,11 +30977,15 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardTimeSeriesRequest.pb( tensorboard_service.GetTensorboardTimeSeriesRequest() ) @@ -30563,6 +31011,10 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.get_tensorboard_time_series( request, @@ -30574,6 +31026,7 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30789,11 +31242,15 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardTimeSeriesRequest.pb( tensorboard_service.UpdateTensorboardTimeSeriesRequest() ) @@ -30819,6 +31276,10 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.update_tensorboard_time_series( request, @@ -30830,6 +31291,7 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30937,11 +31399,15 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardTimeSeriesRequest.pb( tensorboard_service.ListTensorboardTimeSeriesRequest() ) @@ -30967,6 +31433,10 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardTimeSeriesResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardTimeSeriesResponse(), + metadata, + ) await client.list_tensorboard_time_series( request, @@ -30978,6 +31448,7 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31081,11 +31552,15 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardTimeSeriesRequest.pb( tensorboard_service.DeleteTensorboardTimeSeriesRequest() ) @@ -31109,6 +31584,7 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_time_series( request, @@ -31120,6 +31596,7 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31232,11 +31709,15 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors transports.AsyncTensorboardServiceRestInterceptor, "post_batch_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest() ) @@ -31266,6 +31747,10 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors post.return_value = ( tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.batch_read_tensorboard_time_series_data( request, @@ -31277,6 +31762,7 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31387,11 +31873,15 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ReadTensorboardTimeSeriesDataRequest() ) @@ -31419,6 +31909,10 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardTimeSeriesDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.read_tensorboard_time_series_data( request, @@ -31430,6 +31924,7 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31536,11 +32031,15 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_blob_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_blob_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_blob_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardBlobDataRequest.pb( tensorboard_service.ReadTensorboardBlobDataRequest() ) @@ -31566,6 +32065,10 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardBlobDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardBlobDataResponse(), + metadata, + ) await client.read_tensorboard_blob_data( request, @@ -31577,6 +32080,7 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31687,11 +32191,15 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_write_tensorboard_experiment_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_write_tensorboard_experiment_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_write_tensorboard_experiment_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardExperimentDataRequest.pb( tensorboard_service.WriteTensorboardExperimentDataRequest() ) @@ -31719,6 +32227,10 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardExperimentDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardExperimentDataResponse(), + metadata, + ) await client.write_tensorboard_experiment_data( request, @@ -31730,6 +32242,7 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31834,11 +32347,15 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_write_tensorboard_run_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_write_tensorboard_run_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_write_tensorboard_run_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardRunDataRequest.pb( tensorboard_service.WriteTensorboardRunDataRequest() ) @@ -31864,6 +32381,10 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardRunDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardRunDataResponse(), + metadata, + ) await client.write_tensorboard_run_data( request, @@ -31875,6 +32396,7 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31986,11 +32508,15 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_export_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_export_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_export_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ExportTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ExportTensorboardTimeSeriesDataRequest() ) @@ -32020,6 +32546,10 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( post.return_value = ( tensorboard_service.ExportTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.ExportTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.export_tensorboard_time_series_data( request, @@ -32031,6 +32561,7 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_vertex_rag_data_service.py b/tests/unit/gapic/aiplatform_v1/test_vertex_rag_data_service.py index 04b268c6ca..4bd954cb9f 100644 --- a/tests/unit/gapic/aiplatform_v1/test_vertex_rag_data_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_vertex_rag_data_service.py @@ -91,6 +91,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -363,6 +371,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VertexRagDataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VertexRagDataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -7773,10 +7824,14 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_create_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_create_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_create_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.CreateRagCorpusRequest.pb( vertex_rag_data_service.CreateRagCorpusRequest() ) @@ -7800,6 +7855,7 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_rag_corpus( request, @@ -7811,6 +7867,7 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_rag_corpus_rest_bad_request( @@ -7992,10 +8049,14 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_update_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_update_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_update_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UpdateRagCorpusRequest.pb( vertex_rag_data_service.UpdateRagCorpusRequest() ) @@ -8019,6 +8080,7 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_rag_corpus( request, @@ -8030,6 +8092,7 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_rag_corpus_rest_bad_request( @@ -8118,10 +8181,14 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_get_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_get_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_get_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagCorpusRequest.pb( vertex_rag_data_service.GetRagCorpusRequest() ) @@ -8145,6 +8212,7 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagCorpus() + post_with_metadata.return_value = vertex_rag_data.RagCorpus(), metadata client.get_rag_corpus( request, @@ -8156,6 +8224,7 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_rag_corpora_rest_bad_request( @@ -8240,10 +8309,14 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_list_rag_corpora" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_list_rag_corpora_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_list_rag_corpora" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagCorporaRequest.pb( vertex_rag_data_service.ListRagCorporaRequest() ) @@ -8269,6 +8342,10 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagCorporaResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagCorporaResponse(), + metadata, + ) client.list_rag_corpora( request, @@ -8280,6 +8357,7 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_rag_corpus_rest_bad_request( @@ -8360,10 +8438,14 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_delete_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_delete_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_delete_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagCorpusRequest.pb( vertex_rag_data_service.DeleteRagCorpusRequest() ) @@ -8387,6 +8469,7 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_rag_corpus( request, @@ -8398,6 +8481,7 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upload_rag_file_rest_bad_request( @@ -8479,10 +8563,14 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_upload_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_upload_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_upload_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UploadRagFileRequest.pb( vertex_rag_data_service.UploadRagFileRequest() ) @@ -8508,6 +8596,10 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.UploadRagFileResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.UploadRagFileResponse(), + metadata, + ) client.upload_rag_file( request, @@ -8519,6 +8611,7 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_rag_files_rest_bad_request( @@ -8599,10 +8692,14 @@ def test_import_rag_files_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_import_rag_files" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_import_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_import_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ImportRagFilesRequest.pb( vertex_rag_data_service.ImportRagFilesRequest() ) @@ -8626,6 +8723,7 @@ def test_import_rag_files_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_rag_files( request, @@ -8637,6 +8735,7 @@ def test_import_rag_files_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_rag_file_rest_bad_request( @@ -8729,10 +8828,14 @@ def test_get_rag_file_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_get_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_get_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_get_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagFileRequest.pb( vertex_rag_data_service.GetRagFileRequest() ) @@ -8756,6 +8859,7 @@ def test_get_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagFile() + post_with_metadata.return_value = vertex_rag_data.RagFile(), metadata client.get_rag_file( request, @@ -8767,6 +8871,7 @@ def test_get_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_rag_files_rest_bad_request( @@ -8851,10 +8956,14 @@ def test_list_rag_files_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_list_rag_files" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_list_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_list_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagFilesRequest.pb( vertex_rag_data_service.ListRagFilesRequest() ) @@ -8880,6 +8989,10 @@ def test_list_rag_files_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagFilesResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagFilesResponse(), + metadata, + ) client.list_rag_files( request, @@ -8891,6 +9004,7 @@ def test_list_rag_files_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_rag_file_rest_bad_request( @@ -8975,10 +9089,14 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_delete_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_delete_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_delete_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagFileRequest.pb( vertex_rag_data_service.DeleteRagFileRequest() ) @@ -9002,6 +9120,7 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_rag_file( request, @@ -9013,6 +9132,7 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -10072,10 +10192,14 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_create_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_create_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_create_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.CreateRagCorpusRequest.pb( vertex_rag_data_service.CreateRagCorpusRequest() ) @@ -10099,6 +10223,7 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_rag_corpus( request, @@ -10110,6 +10235,7 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10307,10 +10433,14 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_update_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_update_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_update_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UpdateRagCorpusRequest.pb( vertex_rag_data_service.UpdateRagCorpusRequest() ) @@ -10334,6 +10464,7 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_rag_corpus( request, @@ -10345,6 +10476,7 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10449,10 +10581,14 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_get_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_get_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_get_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagCorpusRequest.pb( vertex_rag_data_service.GetRagCorpusRequest() ) @@ -10476,6 +10612,7 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagCorpus() + post_with_metadata.return_value = vertex_rag_data.RagCorpus(), metadata await client.get_rag_corpus( request, @@ -10487,6 +10624,7 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10587,10 +10725,14 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_list_rag_corpora" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_list_rag_corpora_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_list_rag_corpora" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagCorporaRequest.pb( vertex_rag_data_service.ListRagCorporaRequest() ) @@ -10616,6 +10758,10 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagCorporaResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagCorporaResponse(), + metadata, + ) await client.list_rag_corpora( request, @@ -10627,6 +10773,7 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10723,10 +10870,14 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_delete_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_delete_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_delete_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagCorpusRequest.pb( vertex_rag_data_service.DeleteRagCorpusRequest() ) @@ -10750,6 +10901,7 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_rag_corpus( request, @@ -10761,6 +10913,7 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10858,10 +11011,14 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_upload_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_upload_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_upload_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UploadRagFileRequest.pb( vertex_rag_data_service.UploadRagFileRequest() ) @@ -10887,6 +11044,10 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.UploadRagFileResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.UploadRagFileResponse(), + metadata, + ) await client.upload_rag_file( request, @@ -10898,6 +11059,7 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10994,10 +11156,14 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_import_rag_files" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_import_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_import_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ImportRagFilesRequest.pb( vertex_rag_data_service.ImportRagFilesRequest() ) @@ -11021,6 +11187,7 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_rag_files( request, @@ -11032,6 +11199,7 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11140,10 +11308,14 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_get_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_get_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_get_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagFileRequest.pb( vertex_rag_data_service.GetRagFileRequest() ) @@ -11167,6 +11339,7 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagFile() + post_with_metadata.return_value = vertex_rag_data.RagFile(), metadata await client.get_rag_file( request, @@ -11178,6 +11351,7 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11278,10 +11452,14 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_list_rag_files" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_list_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_list_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagFilesRequest.pb( vertex_rag_data_service.ListRagFilesRequest() ) @@ -11307,6 +11485,10 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagFilesResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagFilesResponse(), + metadata, + ) await client.list_rag_files( request, @@ -11318,6 +11500,7 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11418,10 +11601,14 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_delete_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_delete_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_delete_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagFileRequest.pb( vertex_rag_data_service.DeleteRagFileRequest() ) @@ -11445,6 +11632,7 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_rag_file( request, @@ -11456,6 +11644,7 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_vertex_rag_service.py b/tests/unit/gapic/aiplatform_v1/test_vertex_rag_service.py index 255eb71732..bf7c5a8622 100644 --- a/tests/unit/gapic/aiplatform_v1/test_vertex_rag_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_vertex_rag_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -342,6 +350,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VertexRagServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VertexRagServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3097,10 +3148,14 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_retrieve_contexts" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, + "post_retrieve_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_retrieve_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.RetrieveContextsRequest.pb( vertex_rag_service.RetrieveContextsRequest() ) @@ -3126,6 +3181,10 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.RetrieveContextsResponse() + post_with_metadata.return_value = ( + vertex_rag_service.RetrieveContextsResponse(), + metadata, + ) client.retrieve_contexts( request, @@ -3137,6 +3196,7 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_augment_prompt_rest_bad_request( @@ -3218,10 +3278,13 @@ def test_augment_prompt_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_augment_prompt" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, "post_augment_prompt_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_augment_prompt" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.AugmentPromptRequest.pb( vertex_rag_service.AugmentPromptRequest() ) @@ -3247,6 +3310,10 @@ def test_augment_prompt_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.AugmentPromptResponse() + post_with_metadata.return_value = ( + vertex_rag_service.AugmentPromptResponse(), + metadata, + ) client.augment_prompt( request, @@ -3258,6 +3325,7 @@ def test_augment_prompt_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_corroborate_content_rest_bad_request( @@ -3342,10 +3410,14 @@ def test_corroborate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_corroborate_content" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, + "post_corroborate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_corroborate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.CorroborateContentRequest.pb( vertex_rag_service.CorroborateContentRequest() ) @@ -3371,6 +3443,10 @@ def test_corroborate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.CorroborateContentResponse() + post_with_metadata.return_value = ( + vertex_rag_service.CorroborateContentResponse(), + metadata, + ) client.corroborate_content( request, @@ -3382,6 +3458,7 @@ def test_corroborate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -4186,10 +4263,14 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_retrieve_contexts" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_retrieve_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_retrieve_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.RetrieveContextsRequest.pb( vertex_rag_service.RetrieveContextsRequest() ) @@ -4215,6 +4296,10 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.RetrieveContextsResponse() + post_with_metadata.return_value = ( + vertex_rag_service.RetrieveContextsResponse(), + metadata, + ) await client.retrieve_contexts( request, @@ -4226,6 +4311,7 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4323,10 +4409,14 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_augment_prompt" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_augment_prompt_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_augment_prompt" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.AugmentPromptRequest.pb( vertex_rag_service.AugmentPromptRequest() ) @@ -4352,6 +4442,10 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.AugmentPromptResponse() + post_with_metadata.return_value = ( + vertex_rag_service.AugmentPromptResponse(), + metadata, + ) await client.augment_prompt( request, @@ -4363,6 +4457,7 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4463,10 +4558,14 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_corroborate_content" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_corroborate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_corroborate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.CorroborateContentRequest.pb( vertex_rag_service.CorroborateContentRequest() ) @@ -4492,6 +4591,10 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.CorroborateContentResponse() + post_with_metadata.return_value = ( + vertex_rag_service.CorroborateContentResponse(), + metadata, + ) await client.corroborate_content( request, @@ -4503,6 +4606,7 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1/test_vizier_service.py b/tests/unit/gapic/aiplatform_v1/test_vizier_service.py index 09e5ce0ffa..7863d70fc7 100644 --- a/tests/unit/gapic/aiplatform_v1/test_vizier_service.py +++ b/tests/unit/gapic/aiplatform_v1/test_vizier_service.py @@ -86,6 +86,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -339,6 +347,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VizierServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VizierServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -9787,10 +9838,13 @@ def test_create_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_create_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_create_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_create_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateStudyRequest.pb( vizier_service.CreateStudyRequest() ) @@ -9814,6 +9868,7 @@ def test_create_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_study.Study() + post_with_metadata.return_value = gca_study.Study(), metadata client.create_study( request, @@ -9825,6 +9880,7 @@ def test_create_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_study_rest_bad_request(request_type=vizier_service.GetStudyRequest): @@ -9913,10 +9969,13 @@ def test_get_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_get_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_get_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_get_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetStudyRequest.pb(vizier_service.GetStudyRequest()) transcode.return_value = { "method": "post", @@ -9938,6 +9997,7 @@ def test_get_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata client.get_study( request, @@ -9949,6 +10009,7 @@ def test_get_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_studies_rest_bad_request(request_type=vizier_service.ListStudiesRequest): @@ -10031,10 +10092,13 @@ def test_list_studies_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_studies" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_list_studies_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_studies" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListStudiesRequest.pb( vizier_service.ListStudiesRequest() ) @@ -10060,6 +10124,7 @@ def test_list_studies_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListStudiesResponse() + post_with_metadata.return_value = vizier_service.ListStudiesResponse(), metadata client.list_studies( request, @@ -10071,6 +10136,7 @@ def test_list_studies_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_study_rest_bad_request(request_type=vizier_service.DeleteStudyRequest): @@ -10266,10 +10332,13 @@ def test_lookup_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_lookup_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_lookup_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_lookup_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.LookupStudyRequest.pb( vizier_service.LookupStudyRequest() ) @@ -10293,6 +10362,7 @@ def test_lookup_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata client.lookup_study( request, @@ -10304,6 +10374,7 @@ def test_lookup_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_suggest_trials_rest_bad_request( @@ -10384,10 +10455,13 @@ def test_suggest_trials_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VizierServiceRestInterceptor, "post_suggest_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_suggest_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_suggest_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.SuggestTrialsRequest.pb( vizier_service.SuggestTrialsRequest() ) @@ -10411,6 +10485,7 @@ def test_suggest_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.suggest_trials( request, @@ -10422,6 +10497,7 @@ def test_suggest_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_trial_rest_bad_request(request_type=vizier_service.CreateTrialRequest): @@ -10611,10 +10687,13 @@ def test_create_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_create_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_create_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_create_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateTrialRequest.pb( vizier_service.CreateTrialRequest() ) @@ -10638,6 +10717,7 @@ def test_create_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.create_trial( request, @@ -10649,6 +10729,7 @@ def test_create_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_trial_rest_bad_request(request_type=vizier_service.GetTrialRequest): @@ -10745,10 +10826,13 @@ def test_get_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_get_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_get_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_get_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetTrialRequest.pb(vizier_service.GetTrialRequest()) transcode.return_value = { "method": "post", @@ -10770,6 +10854,7 @@ def test_get_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.get_trial( request, @@ -10781,6 +10866,7 @@ def test_get_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_trials_rest_bad_request(request_type=vizier_service.ListTrialsRequest): @@ -10863,10 +10949,13 @@ def test_list_trials_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_list_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListTrialsRequest.pb( vizier_service.ListTrialsRequest() ) @@ -10892,6 +10981,7 @@ def test_list_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListTrialsResponse() + post_with_metadata.return_value = vizier_service.ListTrialsResponse(), metadata client.list_trials( request, @@ -10903,6 +10993,7 @@ def test_list_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_trial_measurement_rest_bad_request( @@ -11001,10 +11092,14 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_add_trial_measurement" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_add_trial_measurement_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_add_trial_measurement" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.AddTrialMeasurementRequest.pb( vizier_service.AddTrialMeasurementRequest() ) @@ -11028,6 +11123,7 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.add_trial_measurement( request, @@ -11039,6 +11135,7 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_complete_trial_rest_bad_request( @@ -11137,10 +11234,13 @@ def test_complete_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_complete_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_complete_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_complete_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CompleteTrialRequest.pb( vizier_service.CompleteTrialRequest() ) @@ -11164,6 +11264,7 @@ def test_complete_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.complete_trial( request, @@ -11175,6 +11276,7 @@ def test_complete_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_trial_rest_bad_request(request_type=vizier_service.DeleteTrialRequest): @@ -11370,10 +11472,14 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VizierServiceRestInterceptor, "post_check_trial_early_stopping_state" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_check_trial_early_stopping_state_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_check_trial_early_stopping_state" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CheckTrialEarlyStoppingStateRequest.pb( vizier_service.CheckTrialEarlyStoppingStateRequest() ) @@ -11397,6 +11503,7 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.check_trial_early_stopping_state( request, @@ -11408,6 +11515,7 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stop_trial_rest_bad_request(request_type=vizier_service.StopTrialRequest): @@ -11504,10 +11612,13 @@ def test_stop_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_stop_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_stop_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_stop_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.StopTrialRequest.pb( vizier_service.StopTrialRequest() ) @@ -11531,6 +11642,7 @@ def test_stop_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.stop_trial( request, @@ -11542,6 +11654,7 @@ def test_stop_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_optimal_trials_rest_bad_request( @@ -11623,10 +11736,14 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_optimal_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_list_optimal_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_optimal_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListOptimalTrialsRequest.pb( vizier_service.ListOptimalTrialsRequest() ) @@ -11652,6 +11769,10 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListOptimalTrialsResponse() + post_with_metadata.return_value = ( + vizier_service.ListOptimalTrialsResponse(), + metadata, + ) client.list_optimal_trials( request, @@ -11663,6 +11784,7 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -12880,10 +13002,13 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_create_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_create_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_create_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateStudyRequest.pb( vizier_service.CreateStudyRequest() ) @@ -12907,6 +13032,7 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_study.Study() + post_with_metadata.return_value = gca_study.Study(), metadata await client.create_study( request, @@ -12918,6 +13044,7 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13024,10 +13151,13 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_get_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_get_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_get_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetStudyRequest.pb(vizier_service.GetStudyRequest()) transcode.return_value = { "method": "post", @@ -13049,6 +13179,7 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata await client.get_study( request, @@ -13060,6 +13191,7 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13160,10 +13292,13 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_studies" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_list_studies_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_studies" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListStudiesRequest.pb( vizier_service.ListStudiesRequest() ) @@ -13189,6 +13324,7 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListStudiesResponse() + post_with_metadata.return_value = vizier_service.ListStudiesResponse(), metadata await client.list_studies( request, @@ -13200,6 +13336,7 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13431,10 +13568,13 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_lookup_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_lookup_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_lookup_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.LookupStudyRequest.pb( vizier_service.LookupStudyRequest() ) @@ -13458,6 +13598,7 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata await client.lookup_study( request, @@ -13469,6 +13610,7 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13565,10 +13707,14 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_suggest_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_suggest_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_suggest_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.SuggestTrialsRequest.pb( vizier_service.SuggestTrialsRequest() ) @@ -13592,6 +13738,7 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.suggest_trials( request, @@ -13603,6 +13750,7 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13810,10 +13958,13 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_create_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_create_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_create_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateTrialRequest.pb( vizier_service.CreateTrialRequest() ) @@ -13837,6 +13988,7 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.create_trial( request, @@ -13848,6 +14000,7 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13962,10 +14115,13 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_get_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_get_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_get_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetTrialRequest.pb(vizier_service.GetTrialRequest()) transcode.return_value = { "method": "post", @@ -13987,6 +14143,7 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.get_trial( request, @@ -13998,6 +14155,7 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14098,10 +14256,13 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_list_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListTrialsRequest.pb( vizier_service.ListTrialsRequest() ) @@ -14127,6 +14288,7 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListTrialsResponse() + post_with_metadata.return_value = vizier_service.ListTrialsResponse(), metadata await client.list_trials( request, @@ -14138,6 +14300,7 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14252,10 +14415,14 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_add_trial_measurement" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_add_trial_measurement_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_add_trial_measurement" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.AddTrialMeasurementRequest.pb( vizier_service.AddTrialMeasurementRequest() ) @@ -14279,6 +14446,7 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.add_trial_measurement( request, @@ -14290,6 +14458,7 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14404,10 +14573,14 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_complete_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_complete_trial_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_complete_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CompleteTrialRequest.pb( vizier_service.CompleteTrialRequest() ) @@ -14431,6 +14604,7 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.complete_trial( request, @@ -14442,6 +14616,7 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14674,11 +14849,15 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( transports.AsyncVizierServiceRestInterceptor, "post_check_trial_early_stopping_state", ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_check_trial_early_stopping_state_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_check_trial_early_stopping_state", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CheckTrialEarlyStoppingStateRequest.pb( vizier_service.CheckTrialEarlyStoppingStateRequest() ) @@ -14702,6 +14881,7 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.check_trial_early_stopping_state( request, @@ -14713,6 +14893,7 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14827,10 +15008,13 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_stop_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_stop_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_stop_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.StopTrialRequest.pb( vizier_service.StopTrialRequest() ) @@ -14854,6 +15038,7 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.stop_trial( request, @@ -14865,6 +15050,7 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14962,10 +15148,14 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_optimal_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_list_optimal_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_optimal_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListOptimalTrialsRequest.pb( vizier_service.ListOptimalTrialsRequest() ) @@ -14991,6 +15181,10 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListOptimalTrialsResponse() + post_with_metadata.return_value = ( + vizier_service.ListOptimalTrialsResponse(), + metadata, + ) await client.list_optimal_trials( request, @@ -15002,6 +15196,7 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_dataset_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_dataset_service.py index 8b49dddeaa..ecdbaaaecf 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_dataset_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_dataset_service.py @@ -99,6 +99,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -356,6 +364,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = DatasetServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = DatasetServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -14086,10 +14137,13 @@ def test_create_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_create_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_create_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_create_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetRequest.pb( dataset_service.CreateDatasetRequest() ) @@ -14113,6 +14167,7 @@ def test_create_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_dataset( request, @@ -14124,6 +14179,7 @@ def test_create_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_dataset_rest_bad_request(request_type=dataset_service.GetDatasetRequest): @@ -14224,10 +14280,13 @@ def test_get_dataset_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_get_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetRequest.pb( dataset_service.GetDatasetRequest() ) @@ -14251,6 +14310,7 @@ def test_get_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset.Dataset() + post_with_metadata.return_value = dataset.Dataset(), metadata client.get_dataset( request, @@ -14262,6 +14322,7 @@ def test_get_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_dataset_rest_bad_request( @@ -14473,10 +14534,13 @@ def test_update_dataset_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_update_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_update_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_update_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetRequest.pb( dataset_service.UpdateDatasetRequest() ) @@ -14500,6 +14564,7 @@ def test_update_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset.Dataset() + post_with_metadata.return_value = gca_dataset.Dataset(), metadata client.update_dataset( request, @@ -14511,6 +14576,7 @@ def test_update_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_datasets_rest_bad_request( @@ -14595,10 +14661,13 @@ def test_list_datasets_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_datasets" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_datasets_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_datasets" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetsRequest.pb( dataset_service.ListDatasetsRequest() ) @@ -14624,6 +14693,10 @@ def test_list_datasets_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetsResponse(), + metadata, + ) client.list_datasets( request, @@ -14635,6 +14708,7 @@ def test_list_datasets_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_dataset_rest_bad_request( @@ -14715,10 +14789,13 @@ def test_delete_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_dataset" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_delete_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetRequest.pb( dataset_service.DeleteDatasetRequest() ) @@ -14742,6 +14819,7 @@ def test_delete_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_dataset( request, @@ -14753,6 +14831,7 @@ def test_delete_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_data_rest_bad_request(request_type=dataset_service.ImportDataRequest): @@ -14831,10 +14910,13 @@ def test_import_data_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_import_data" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_import_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_import_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ImportDataRequest.pb( dataset_service.ImportDataRequest() ) @@ -14858,6 +14940,7 @@ def test_import_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_data( request, @@ -14869,6 +14952,7 @@ def test_import_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_data_rest_bad_request(request_type=dataset_service.ExportDataRequest): @@ -14947,10 +15031,13 @@ def test_export_data_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_export_data" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_export_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_export_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ExportDataRequest.pb( dataset_service.ExportDataRequest() ) @@ -14974,6 +15061,7 @@ def test_export_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_data( request, @@ -14985,6 +15073,7 @@ def test_export_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_dataset_version_rest_bad_request( @@ -15153,10 +15242,14 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_create_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_create_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_create_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetVersionRequest.pb( dataset_service.CreateDatasetVersionRequest() ) @@ -15180,6 +15273,7 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_dataset_version( request, @@ -15191,6 +15285,7 @@ def test_create_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_dataset_version_rest_bad_request( @@ -15383,10 +15478,14 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_update_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_update_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_update_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetVersionRequest.pb( dataset_service.UpdateDatasetVersionRequest() ) @@ -15412,6 +15511,7 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset_version.DatasetVersion() + post_with_metadata.return_value = gca_dataset_version.DatasetVersion(), metadata client.update_dataset_version( request, @@ -15423,6 +15523,7 @@ def test_update_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_dataset_version_rest_bad_request( @@ -15507,10 +15608,14 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_delete_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetVersionRequest.pb( dataset_service.DeleteDatasetVersionRequest() ) @@ -15534,6 +15639,7 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_dataset_version( request, @@ -15545,6 +15651,7 @@ def test_delete_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_dataset_version_rest_bad_request( @@ -15645,10 +15752,14 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_get_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetVersionRequest.pb( dataset_service.GetDatasetVersionRequest() ) @@ -15674,6 +15785,7 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_version.DatasetVersion() + post_with_metadata.return_value = dataset_version.DatasetVersion(), metadata client.get_dataset_version( request, @@ -15685,6 +15797,7 @@ def test_get_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_dataset_versions_rest_bad_request( @@ -15769,10 +15882,14 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_dataset_versions" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_list_dataset_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_dataset_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetVersionsRequest.pb( dataset_service.ListDatasetVersionsRequest() ) @@ -15798,6 +15915,10 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetVersionsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetVersionsResponse(), + metadata, + ) client.list_dataset_versions( request, @@ -15809,6 +15930,7 @@ def test_list_dataset_versions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_restore_dataset_version_rest_bad_request( @@ -15893,10 +16015,14 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_restore_dataset_version" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_restore_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_restore_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.RestoreDatasetVersionRequest.pb( dataset_service.RestoreDatasetVersionRequest() ) @@ -15920,6 +16046,7 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.restore_dataset_version( request, @@ -15931,6 +16058,7 @@ def test_restore_dataset_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_data_items_rest_bad_request( @@ -16015,10 +16143,13 @@ def test_list_data_items_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_data_items" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_data_items_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDataItemsRequest.pb( dataset_service.ListDataItemsRequest() ) @@ -16044,6 +16175,10 @@ def test_list_data_items_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDataItemsResponse(), + metadata, + ) client.list_data_items( request, @@ -16055,6 +16190,7 @@ def test_list_data_items_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_data_items_rest_bad_request( @@ -16139,10 +16275,13 @@ def test_search_data_items_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_search_data_items" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_search_data_items_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_search_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.SearchDataItemsRequest.pb( dataset_service.SearchDataItemsRequest() ) @@ -16168,6 +16307,10 @@ def test_search_data_items_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.SearchDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.SearchDataItemsResponse(), + metadata, + ) client.search_data_items( request, @@ -16179,6 +16322,7 @@ def test_search_data_items_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_saved_queries_rest_bad_request( @@ -16263,10 +16407,14 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_saved_queries" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_list_saved_queries_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_saved_queries" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListSavedQueriesRequest.pb( dataset_service.ListSavedQueriesRequest() ) @@ -16292,6 +16440,10 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListSavedQueriesResponse() + post_with_metadata.return_value = ( + dataset_service.ListSavedQueriesResponse(), + metadata, + ) client.list_saved_queries( request, @@ -16303,6 +16455,7 @@ def test_list_saved_queries_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_saved_query_rest_bad_request( @@ -16387,10 +16540,14 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): ), mock.patch.object( transports.DatasetServiceRestInterceptor, "post_delete_saved_query" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_delete_saved_query_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_delete_saved_query" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteSavedQueryRequest.pb( dataset_service.DeleteSavedQueryRequest() ) @@ -16414,6 +16571,7 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_saved_query( request, @@ -16425,6 +16583,7 @@ def test_delete_saved_query_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_annotation_spec_rest_bad_request( @@ -16517,10 +16676,14 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_get_annotation_spec" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, + "post_get_annotation_spec_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_get_annotation_spec" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetAnnotationSpecRequest.pb( dataset_service.GetAnnotationSpecRequest() ) @@ -16546,6 +16709,7 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = annotation_spec.AnnotationSpec() + post_with_metadata.return_value = annotation_spec.AnnotationSpec(), metadata client.get_annotation_spec( request, @@ -16557,6 +16721,7 @@ def test_get_annotation_spec_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_annotations_rest_bad_request( @@ -16645,10 +16810,13 @@ def test_list_annotations_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.DatasetServiceRestInterceptor, "post_list_annotations" ) as post, mock.patch.object( + transports.DatasetServiceRestInterceptor, "post_list_annotations_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.DatasetServiceRestInterceptor, "pre_list_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListAnnotationsRequest.pb( dataset_service.ListAnnotationsRequest() ) @@ -16674,6 +16842,10 @@ def test_list_annotations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListAnnotationsResponse() + post_with_metadata.return_value = ( + dataset_service.ListAnnotationsResponse(), + metadata, + ) client.list_annotations( request, @@ -16685,6 +16857,7 @@ def test_list_annotations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -17946,10 +18119,14 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_create_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_create_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_create_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetRequest.pb( dataset_service.CreateDatasetRequest() ) @@ -17973,6 +18150,7 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_dataset( request, @@ -17984,6 +18162,7 @@ async def test_create_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18102,10 +18281,13 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetRequest.pb( dataset_service.GetDatasetRequest() ) @@ -18129,6 +18311,7 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset.Dataset() + post_with_metadata.return_value = dataset.Dataset(), metadata await client.get_dataset( request, @@ -18140,6 +18323,7 @@ async def test_get_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18367,10 +18551,14 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_update_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_update_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_update_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetRequest.pb( dataset_service.UpdateDatasetRequest() ) @@ -18394,6 +18582,7 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_dataset.Dataset() + post_with_metadata.return_value = gca_dataset.Dataset(), metadata await client.update_dataset( request, @@ -18405,6 +18594,7 @@ async def test_update_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18505,10 +18695,14 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_datasets" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_datasets_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_datasets" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetsRequest.pb( dataset_service.ListDatasetsRequest() ) @@ -18534,6 +18728,10 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetsResponse(), + metadata, + ) await client.list_datasets( request, @@ -18545,6 +18743,7 @@ async def test_list_datasets_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18641,10 +18840,14 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_dataset" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetRequest.pb( dataset_service.DeleteDatasetRequest() ) @@ -18668,6 +18871,7 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_dataset( request, @@ -18679,6 +18883,7 @@ async def test_delete_dataset_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18775,10 +18980,13 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_import_data" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_import_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_import_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ImportDataRequest.pb( dataset_service.ImportDataRequest() ) @@ -18802,6 +19010,7 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_data( request, @@ -18813,6 +19022,7 @@ async def test_import_data_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18909,10 +19119,13 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_export_data" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, "post_export_data_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_export_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ExportDataRequest.pb( dataset_service.ExportDataRequest() ) @@ -18936,6 +19149,7 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_data( request, @@ -18947,6 +19161,7 @@ async def test_export_data_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19131,10 +19346,14 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_create_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_create_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_create_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.CreateDatasetVersionRequest.pb( dataset_service.CreateDatasetVersionRequest() ) @@ -19158,6 +19377,7 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_dataset_version( request, @@ -19169,6 +19389,7 @@ async def test_create_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19377,10 +19598,14 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_update_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_update_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_update_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.UpdateDatasetVersionRequest.pb( dataset_service.UpdateDatasetVersionRequest() ) @@ -19406,6 +19631,7 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_dataset_version.DatasetVersion() + post_with_metadata.return_value = gca_dataset_version.DatasetVersion(), metadata await client.update_dataset_version( request, @@ -19417,6 +19643,7 @@ async def test_update_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19517,10 +19744,14 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteDatasetVersionRequest.pb( dataset_service.DeleteDatasetVersionRequest() ) @@ -19544,6 +19775,7 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_dataset_version( request, @@ -19555,6 +19787,7 @@ async def test_delete_dataset_version_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19671,10 +19904,14 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_get_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetDatasetVersionRequest.pb( dataset_service.GetDatasetVersionRequest() ) @@ -19700,6 +19937,7 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_version.DatasetVersion() + post_with_metadata.return_value = dataset_version.DatasetVersion(), metadata await client.get_dataset_version( request, @@ -19711,6 +19949,7 @@ async def test_get_dataset_version_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19811,10 +20050,14 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_dataset_versions" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_dataset_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_dataset_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDatasetVersionsRequest.pb( dataset_service.ListDatasetVersionsRequest() ) @@ -19840,6 +20083,10 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = dataset_service.ListDatasetVersionsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDatasetVersionsResponse(), + metadata, + ) await client.list_dataset_versions( request, @@ -19851,6 +20098,7 @@ async def test_list_dataset_versions_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19951,10 +20199,14 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_restore_dataset_version" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_restore_dataset_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_restore_dataset_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.RestoreDatasetVersionRequest.pb( dataset_service.RestoreDatasetVersionRequest() ) @@ -19978,6 +20230,7 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.restore_dataset_version( request, @@ -19989,6 +20242,7 @@ async def test_restore_dataset_version_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20089,10 +20343,14 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_data_items" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_data_items_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListDataItemsRequest.pb( dataset_service.ListDataItemsRequest() ) @@ -20118,6 +20376,10 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.ListDataItemsResponse(), + metadata, + ) await client.list_data_items( request, @@ -20129,6 +20391,7 @@ async def test_list_data_items_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20229,10 +20492,14 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_search_data_items" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_search_data_items_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_search_data_items" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.SearchDataItemsRequest.pb( dataset_service.SearchDataItemsRequest() ) @@ -20258,6 +20525,10 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.SearchDataItemsResponse() + post_with_metadata.return_value = ( + dataset_service.SearchDataItemsResponse(), + metadata, + ) await client.search_data_items( request, @@ -20269,6 +20540,7 @@ async def test_search_data_items_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20369,10 +20641,14 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_saved_queries" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_saved_queries_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_saved_queries" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListSavedQueriesRequest.pb( dataset_service.ListSavedQueriesRequest() ) @@ -20398,6 +20674,10 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListSavedQueriesResponse() + post_with_metadata.return_value = ( + dataset_service.ListSavedQueriesResponse(), + metadata, + ) await client.list_saved_queries( request, @@ -20409,6 +20689,7 @@ async def test_list_saved_queries_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20509,10 +20790,14 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_delete_saved_query" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_delete_saved_query_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_delete_saved_query" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.DeleteSavedQueryRequest.pb( dataset_service.DeleteSavedQueryRequest() ) @@ -20536,6 +20821,7 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_saved_query( request, @@ -20547,6 +20833,7 @@ async def test_delete_saved_query_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20655,10 +20942,14 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_get_annotation_spec" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_get_annotation_spec_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_get_annotation_spec" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.GetAnnotationSpecRequest.pb( dataset_service.GetAnnotationSpecRequest() ) @@ -20684,6 +20975,7 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = annotation_spec.AnnotationSpec() + post_with_metadata.return_value = annotation_spec.AnnotationSpec(), metadata await client.get_annotation_spec( request, @@ -20695,6 +20987,7 @@ async def test_get_annotation_spec_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20799,10 +21092,14 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "post_list_annotations" ) as post, mock.patch.object( + transports.AsyncDatasetServiceRestInterceptor, + "post_list_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDatasetServiceRestInterceptor, "pre_list_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = dataset_service.ListAnnotationsRequest.pb( dataset_service.ListAnnotationsRequest() ) @@ -20828,6 +21125,10 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = dataset_service.ListAnnotationsResponse() + post_with_metadata.return_value = ( + dataset_service.ListAnnotationsResponse(), + metadata, + ) await client.list_annotations( request, @@ -20839,6 +21140,7 @@ async def test_list_annotations_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_deployment_resource_pool_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_deployment_resource_pool_service.py index 52534574e0..dd913f3cdd 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_deployment_resource_pool_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_deployment_resource_pool_service.py @@ -101,6 +101,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -387,6 +395,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = DeploymentResourcePoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = DeploymentResourcePoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5658,11 +5709,15 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_create_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_create_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_create_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest() @@ -5688,6 +5743,7 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_deployment_resource_pool( request, @@ -5699,6 +5755,7 @@ def test_create_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_deployment_resource_pool_rest_bad_request( @@ -5796,11 +5853,15 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_get_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_get_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_get_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.GetDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.GetDeploymentResourcePoolRequest() @@ -5828,6 +5889,10 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = deployment_resource_pool.DeploymentResourcePool() + post_with_metadata.return_value = ( + deployment_resource_pool.DeploymentResourcePool(), + metadata, + ) client.get_deployment_resource_pool( request, @@ -5839,6 +5904,7 @@ def test_get_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_deployment_resource_pools_rest_bad_request( @@ -5930,11 +5996,15 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_list_deployment_resource_pools", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_list_deployment_resource_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_list_deployment_resource_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest.pb( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest() @@ -5964,6 +6034,10 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): post.return_value = ( deployment_resource_pool_service.ListDeploymentResourcePoolsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse(), + metadata, + ) client.list_deployment_resource_pools( request, @@ -5975,6 +6049,7 @@ def test_list_deployment_resource_pools_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_deployment_resource_pool_rest_bad_request( @@ -6164,11 +6239,15 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_update_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_update_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_update_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest() @@ -6194,6 +6273,7 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_deployment_resource_pool( request, @@ -6205,6 +6285,7 @@ def test_update_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_deployment_resource_pool_rest_bad_request( @@ -6290,11 +6371,15 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_delete_deployment_resource_pool", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_delete_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_delete_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest() @@ -6320,6 +6405,7 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_deployment_resource_pool( request, @@ -6331,6 +6417,7 @@ def test_delete_deployment_resource_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_deployed_models_rest_bad_request( @@ -6426,11 +6513,15 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): transports.DeploymentResourcePoolServiceRestInterceptor, "post_query_deployed_models", ) as post, mock.patch.object( + transports.DeploymentResourcePoolServiceRestInterceptor, + "post_query_deployed_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.DeploymentResourcePoolServiceRestInterceptor, "pre_query_deployed_models", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = deployment_resource_pool_service.QueryDeployedModelsRequest.pb( deployment_resource_pool_service.QueryDeployedModelsRequest() ) @@ -6460,6 +6551,10 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): post.return_value = ( deployment_resource_pool_service.QueryDeployedModelsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.QueryDeployedModelsResponse(), + metadata, + ) client.query_deployed_models( request, @@ -6471,6 +6566,7 @@ def test_query_deployed_models_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7372,11 +7468,15 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_create_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_create_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_create_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.CreateDeploymentResourcePoolRequest() @@ -7402,6 +7502,7 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_deployment_resource_pool( request, @@ -7413,6 +7514,7 @@ async def test_create_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7526,11 +7628,15 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_get_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_get_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_get_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.GetDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.GetDeploymentResourcePoolRequest() @@ -7558,6 +7664,10 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = deployment_resource_pool.DeploymentResourcePool() + post_with_metadata.return_value = ( + deployment_resource_pool.DeploymentResourcePool(), + metadata, + ) await client.get_deployment_resource_pool( request, @@ -7569,6 +7679,7 @@ async def test_get_deployment_resource_pool_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7678,11 +7789,15 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_list_deployment_resource_pools", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_list_deployment_resource_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_list_deployment_resource_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest.pb( deployment_resource_pool_service.ListDeploymentResourcePoolsRequest() @@ -7712,6 +7827,10 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( post.return_value = ( deployment_resource_pool_service.ListDeploymentResourcePoolsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.ListDeploymentResourcePoolsResponse(), + metadata, + ) await client.list_deployment_resource_pools( request, @@ -7723,6 +7842,7 @@ async def test_list_deployment_resource_pools_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7930,11 +8050,15 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_update_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_update_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_update_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.UpdateDeploymentResourcePoolRequest() @@ -7960,6 +8084,7 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_deployment_resource_pool( request, @@ -7971,6 +8096,7 @@ async def test_update_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8074,11 +8200,15 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_delete_deployment_resource_pool", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_delete_deployment_resource_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_delete_deployment_resource_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest.pb( deployment_resource_pool_service.DeleteDeploymentResourcePoolRequest() @@ -8104,6 +8234,7 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_deployment_resource_pool( request, @@ -8115,6 +8246,7 @@ async def test_delete_deployment_resource_pool_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8226,11 +8358,15 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "post_query_deployed_models", ) as post, mock.patch.object( + transports.AsyncDeploymentResourcePoolServiceRestInterceptor, + "post_query_deployed_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncDeploymentResourcePoolServiceRestInterceptor, "pre_query_deployed_models", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = deployment_resource_pool_service.QueryDeployedModelsRequest.pb( deployment_resource_pool_service.QueryDeployedModelsRequest() ) @@ -8260,6 +8396,10 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) post.return_value = ( deployment_resource_pool_service.QueryDeployedModelsResponse() ) + post_with_metadata.return_value = ( + deployment_resource_pool_service.QueryDeployedModelsResponse(), + metadata, + ) await client.query_deployed_models( request, @@ -8271,6 +8411,7 @@ async def test_query_deployed_models_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_endpoint_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_endpoint_service.py index 8ed0b94fe0..bf7ba70ef1 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_endpoint_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_endpoint_service.py @@ -100,6 +100,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -357,6 +365,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = EndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = EndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -7036,6 +7087,14 @@ def test_create_endpoint_rest_call_success(request_type): "service_attachment": "service_attachment_value", }, "faster_deployment_config": {"fast_tryout_enabled": True}, + "rollout_options": { + "max_unavailable_replicas": 2523, + "max_unavailable_percentage": 2726, + "max_surge_replicas": 1917, + "max_surge_percentage": 2120, + "previous_deployed_model": "previous_deployed_model_value", + "revision_number": 1623, + }, "status": { "message": "message_value", "last_update_time": {}, @@ -7181,10 +7240,13 @@ def test_create_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_create_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_create_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_create_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.CreateEndpointRequest.pb( endpoint_service.CreateEndpointRequest() ) @@ -7208,6 +7270,7 @@ def test_create_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_endpoint( request, @@ -7219,6 +7282,7 @@ def test_create_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_endpoint_rest_bad_request( @@ -7326,10 +7390,13 @@ def test_get_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_get_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_get_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_get_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.GetEndpointRequest.pb( endpoint_service.GetEndpointRequest() ) @@ -7353,6 +7420,7 @@ def test_get_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint.Endpoint() + post_with_metadata.return_value = endpoint.Endpoint(), metadata client.get_endpoint( request, @@ -7364,6 +7432,7 @@ def test_get_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_endpoints_rest_bad_request( @@ -7448,10 +7517,13 @@ def test_list_endpoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_list_endpoints" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_list_endpoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_list_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.ListEndpointsRequest.pb( endpoint_service.ListEndpointsRequest() ) @@ -7477,6 +7549,10 @@ def test_list_endpoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint_service.ListEndpointsResponse() + post_with_metadata.return_value = ( + endpoint_service.ListEndpointsResponse(), + metadata, + ) client.list_endpoints( request, @@ -7488,6 +7564,7 @@ def test_list_endpoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_endpoint_rest_bad_request( @@ -7628,6 +7705,14 @@ def test_update_endpoint_rest_call_success(request_type): "service_attachment": "service_attachment_value", }, "faster_deployment_config": {"fast_tryout_enabled": True}, + "rollout_options": { + "max_unavailable_replicas": 2523, + "max_unavailable_percentage": 2726, + "max_surge_replicas": 1917, + "max_surge_percentage": 2120, + "previous_deployed_model": "previous_deployed_model_value", + "revision_number": 1623, + }, "status": { "message": "message_value", "last_update_time": {}, @@ -7800,10 +7885,13 @@ def test_update_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.EndpointServiceRestInterceptor, "post_update_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_update_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_update_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointRequest.pb( endpoint_service.UpdateEndpointRequest() ) @@ -7827,6 +7915,7 @@ def test_update_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_endpoint.Endpoint() + post_with_metadata.return_value = gca_endpoint.Endpoint(), metadata client.update_endpoint( request, @@ -7838,6 +7927,7 @@ def test_update_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_endpoint_long_running_rest_bad_request( @@ -7922,10 +8012,14 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_update_endpoint_long_running" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, + "post_update_endpoint_long_running_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_update_endpoint_long_running" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointLongRunningRequest.pb( endpoint_service.UpdateEndpointLongRunningRequest() ) @@ -7949,6 +8043,7 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_endpoint_long_running( request, @@ -7960,6 +8055,7 @@ def test_update_endpoint_long_running_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_endpoint_rest_bad_request( @@ -8040,10 +8136,13 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_delete_endpoint" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_delete_endpoint_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_delete_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeleteEndpointRequest.pb( endpoint_service.DeleteEndpointRequest() ) @@ -8067,6 +8166,7 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_endpoint( request, @@ -8078,6 +8178,7 @@ def test_delete_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_deploy_model_rest_bad_request( @@ -8158,10 +8259,13 @@ def test_deploy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_deploy_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_deploy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_deploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeployModelRequest.pb( endpoint_service.DeployModelRequest() ) @@ -8185,6 +8289,7 @@ def test_deploy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.deploy_model( request, @@ -8196,6 +8301,7 @@ def test_deploy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_undeploy_model_rest_bad_request( @@ -8276,10 +8382,13 @@ def test_undeploy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_undeploy_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, "post_undeploy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_undeploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UndeployModelRequest.pb( endpoint_service.UndeployModelRequest() ) @@ -8303,6 +8412,7 @@ def test_undeploy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.undeploy_model( request, @@ -8314,6 +8424,7 @@ def test_undeploy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_mutate_deployed_model_rest_bad_request( @@ -8394,10 +8505,14 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.EndpointServiceRestInterceptor, "post_mutate_deployed_model" ) as post, mock.patch.object( + transports.EndpointServiceRestInterceptor, + "post_mutate_deployed_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.EndpointServiceRestInterceptor, "pre_mutate_deployed_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.MutateDeployedModelRequest.pb( endpoint_service.MutateDeployedModelRequest() ) @@ -8421,6 +8536,7 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.mutate_deployed_model( request, @@ -8432,6 +8548,7 @@ def test_mutate_deployed_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -9421,6 +9538,14 @@ async def test_create_endpoint_rest_asyncio_call_success(request_type): "service_attachment": "service_attachment_value", }, "faster_deployment_config": {"fast_tryout_enabled": True}, + "rollout_options": { + "max_unavailable_replicas": 2523, + "max_unavailable_percentage": 2726, + "max_surge_replicas": 1917, + "max_surge_percentage": 2120, + "previous_deployed_model": "previous_deployed_model_value", + "revision_number": 1623, + }, "status": { "message": "message_value", "last_update_time": {}, @@ -9573,10 +9698,14 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_create_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_create_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_create_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.CreateEndpointRequest.pb( endpoint_service.CreateEndpointRequest() ) @@ -9600,6 +9729,7 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_endpoint( request, @@ -9611,6 +9741,7 @@ async def test_create_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9734,10 +9865,14 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_get_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_get_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_get_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.GetEndpointRequest.pb( endpoint_service.GetEndpointRequest() ) @@ -9761,6 +9896,7 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint.Endpoint() + post_with_metadata.return_value = endpoint.Endpoint(), metadata await client.get_endpoint( request, @@ -9772,6 +9908,7 @@ async def test_get_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9872,10 +10009,14 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_list_endpoints" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_list_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_list_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.ListEndpointsRequest.pb( endpoint_service.ListEndpointsRequest() ) @@ -9901,6 +10042,10 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = endpoint_service.ListEndpointsResponse() + post_with_metadata.return_value = ( + endpoint_service.ListEndpointsResponse(), + metadata, + ) await client.list_endpoints( request, @@ -9912,6 +10057,7 @@ async def test_list_endpoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10061,6 +10207,14 @@ async def test_update_endpoint_rest_asyncio_call_success(request_type): "service_attachment": "service_attachment_value", }, "faster_deployment_config": {"fast_tryout_enabled": True}, + "rollout_options": { + "max_unavailable_replicas": 2523, + "max_unavailable_percentage": 2726, + "max_surge_replicas": 1917, + "max_surge_percentage": 2120, + "previous_deployed_model": "previous_deployed_model_value", + "revision_number": 1623, + }, "status": { "message": "message_value", "last_update_time": {}, @@ -10240,10 +10394,14 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_update_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_update_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_update_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointRequest.pb( endpoint_service.UpdateEndpointRequest() ) @@ -10267,6 +10425,7 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_endpoint.Endpoint() + post_with_metadata.return_value = gca_endpoint.Endpoint(), metadata await client.update_endpoint( request, @@ -10278,6 +10437,7 @@ async def test_update_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10379,11 +10539,15 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter transports.AsyncEndpointServiceRestInterceptor, "post_update_endpoint_long_running", ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_update_endpoint_long_running_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_update_endpoint_long_running", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UpdateEndpointLongRunningRequest.pb( endpoint_service.UpdateEndpointLongRunningRequest() ) @@ -10407,6 +10571,7 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_endpoint_long_running( request, @@ -10418,6 +10583,7 @@ async def test_update_endpoint_long_running_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10514,10 +10680,14 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_delete_endpoint" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_delete_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_delete_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeleteEndpointRequest.pb( endpoint_service.DeleteEndpointRequest() ) @@ -10541,6 +10711,7 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_endpoint( request, @@ -10552,6 +10723,7 @@ async def test_delete_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10648,10 +10820,14 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_deploy_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_deploy_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_deploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.DeployModelRequest.pb( endpoint_service.DeployModelRequest() ) @@ -10675,6 +10851,7 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.deploy_model( request, @@ -10686,6 +10863,7 @@ async def test_deploy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10782,10 +10960,14 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_undeploy_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_undeploy_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_undeploy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.UndeployModelRequest.pb( endpoint_service.UndeployModelRequest() ) @@ -10809,6 +10991,7 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.undeploy_model( request, @@ -10820,6 +11003,7 @@ async def test_undeploy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10916,10 +11100,14 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "post_mutate_deployed_model" ) as post, mock.patch.object( + transports.AsyncEndpointServiceRestInterceptor, + "post_mutate_deployed_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEndpointServiceRestInterceptor, "pre_mutate_deployed_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = endpoint_service.MutateDeployedModelRequest.pb( endpoint_service.MutateDeployedModelRequest() ) @@ -10943,6 +11131,7 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.mutate_deployed_model( request, @@ -10954,6 +11143,7 @@ async def test_mutate_deployed_model_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_evaluation_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_evaluation_service.py index 53a6ba23f7..de907c96c2 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_evaluation_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_evaluation_service.py @@ -55,9 +55,13 @@ from google.api_core import client_options from google.api_core import exceptions as core_exceptions +from google.api_core import future from google.api_core import gapic_v1 from google.api_core import grpc_helpers from google.api_core import grpc_helpers_async +from google.api_core import operation +from google.api_core import operation_async # type: ignore +from google.api_core import operations_v1 from google.api_core import path_template from google.api_core import retry as retries from google.auth import credentials as ga_credentials @@ -70,6 +74,7 @@ ) from google.cloud.aiplatform_v1beta1.services.evaluation_service import transports from google.cloud.aiplatform_v1beta1.types import evaluation_service +from google.cloud.aiplatform_v1beta1.types import io from google.cloud.location import locations_pb2 from google.iam.v1 import iam_policy_pb2 # type: ignore from google.iam.v1 import options_pb2 # type: ignore @@ -79,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -345,6 +358,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = EvaluationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = EvaluationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1366,6 +1422,254 @@ async def test_evaluate_instances_field_headers_async(): ) in kw["metadata"] +@pytest.mark.parametrize( + "request_type", + [ + evaluation_service.EvaluateDatasetRequest, + dict, + ], +) +def test_evaluate_dataset(request_type, transport: str = "grpc"): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/spam") + response = client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = evaluation_service.EvaluateDatasetRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +def test_evaluate_dataset_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are + # automatically populated, according to AIP-4235, with non-empty requests. + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Populate all string fields in the request which are not UUID4 + # since we want to check that UUID4 are populated automatically + # if they meet the requirements of AIP 4235. + request = evaluation_service.EvaluateDatasetRequest( + location="location_value", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + call.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client.evaluate_dataset(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == evaluation_service.EvaluateDatasetRequest( + location="location_value", + ) + + +def test_evaluate_dataset_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.evaluate_dataset in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[ + client._transport.evaluate_dataset + ] = mock_rpc + request = {} + client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + # Operation methods call wrapper_fn to build a cached + # client._transport.operations_client instance on first rpc call. + # Subsequent calls should use the cached wrapper + wrapper_fn.reset_mock() + + client.evaluate_dataset(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_evaluate_dataset_async_use_cached_wrapped_rpc( + transport: str = "grpc_asyncio", +): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method_async.wrap_method") as wrapper_fn: + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert ( + client._client._transport.evaluate_dataset + in client._client._transport._wrapped_methods + ) + + # Replace cached wrapped function with mock + mock_rpc = mock.AsyncMock() + mock_rpc.return_value = mock.Mock() + client._client._transport._wrapped_methods[ + client._client._transport.evaluate_dataset + ] = mock_rpc + + request = {} + await client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + # Operation methods call wrapper_fn to build a cached + # client._transport.operations_client instance on first rpc call. + # Subsequent calls should use the cached wrapper + wrapper_fn.reset_mock() + + await client.evaluate_dataset(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_evaluate_dataset_async( + transport: str = "grpc_asyncio", + request_type=evaluation_service.EvaluateDatasetRequest, +): + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/spam") + ) + response = await client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = evaluation_service.EvaluateDatasetRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +@pytest.mark.asyncio +async def test_evaluate_dataset_async_from_dict(): + await test_evaluate_dataset_async(request_type=dict) + + +def test_evaluate_dataset_field_headers(): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = evaluation_service.EvaluateDatasetRequest() + + request.location = "location_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + call.return_value = operations_pb2.Operation(name="operations/op") + client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "location=location_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_evaluate_dataset_field_headers_async(): + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = evaluation_service.EvaluateDatasetRequest() + + request.location = "location_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/op") + ) + await client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "location=location_value", + ) in kw["metadata"] + + def test_evaluate_instances_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call @@ -1491,6 +1795,140 @@ def test_evaluate_instances_rest_unset_required_fields(): assert set(unset_fields) == (set(()) & set(("location",))) +def test_evaluate_dataset_rest_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.evaluate_dataset in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[ + client._transport.evaluate_dataset + ] = mock_rpc + + request = {} + client.evaluate_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + # Operation methods build a cached wrapper on first rpc call + # subsequent calls should use the cached wrapper + wrapper_fn.reset_mock() + + client.evaluate_dataset(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +def test_evaluate_dataset_rest_required_fields( + request_type=evaluation_service.EvaluateDatasetRequest, +): + transport_class = transports.EvaluationServiceRestTransport + + request_init = {} + request_init["location"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson(pb_request, use_integers_for_enums=False) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).evaluate_dataset._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["location"] = "location_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).evaluate_dataset._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "location" in jsonified_request + assert jsonified_request["location"] == "location_value" + + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "post", + "query_params": pb_request, + } + transcode_result["body"] = pb_request + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + + response = client.evaluate_dataset(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_evaluate_dataset_rest_unset_required_fields(): + transport = transports.EvaluationServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.evaluate_dataset._get_unset_required_fields({}) + assert set(unset_fields) == ( + set(()) + & set( + ( + "location", + "dataset", + "metrics", + "outputConfig", + ) + ) + ) + + def test_credentials_transport_error(): # It is an error to provide credentials and a transport instance. transport = transports.EvaluationServiceGrpcTransport( @@ -1620,6 +2058,27 @@ def test_evaluate_instances_empty_call_grpc(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_evaluate_dataset_empty_call_grpc(): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + call.return_value = operations_pb2.Operation(name="operations/op") + client.evaluate_dataset(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = evaluation_service.EvaluateDatasetRequest() + + assert args[0] == request_msg + + def test_transport_kind_grpc_asyncio(): transport = EvaluationServiceAsyncClient.get_transport_class("grpc_asyncio")( credentials=async_anonymous_credentials() @@ -1651,25 +2110,180 @@ async def test_evaluate_instances_empty_call_grpc_asyncio(): call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( evaluation_service.EvaluateInstancesResponse() ) - await client.evaluate_instances(request=None) - - # Establish that the underlying stub method was called. - call.assert_called() - _, args, _ = call.mock_calls[0] - request_msg = evaluation_service.EvaluateInstancesRequest() + await client.evaluate_instances(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = evaluation_service.EvaluateInstancesRequest() + + assert args[0] == request_msg + + +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_evaluate_dataset_empty_call_grpc_asyncio(): + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/spam") + ) + await client.evaluate_dataset(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = evaluation_service.EvaluateDatasetRequest() + + assert args[0] == request_msg + + +def test_transport_kind_rest(): + transport = EvaluationServiceClient.get_transport_class("rest")( + credentials=ga_credentials.AnonymousCredentials() + ) + assert transport.kind == "rest" + + +def test_evaluate_instances_rest_bad_request( + request_type=evaluation_service.EvaluateInstancesRequest, +): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + # send a request that will satisfy transcoding + request_init = {"location": "projects/sample1/locations/sample2"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + json_return_value = "" + response_value.json = mock.Mock(return_value={}) + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + client.evaluate_instances(request) + + +@pytest.mark.parametrize( + "request_type", + [ + evaluation_service.EvaluateInstancesRequest, + dict, + ], +) +def test_evaluate_instances_rest_call_success(request_type): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + # send a request that will satisfy transcoding + request_init = {"location": "projects/sample1/locations/sample2"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = evaluation_service.EvaluateInstancesResponse() + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = evaluation_service.EvaluateInstancesResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + response_value.content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = client.evaluate_instances(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, evaluation_service.EvaluateInstancesResponse) + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_evaluate_instances_rest_interceptors(null_interceptor): + transport = transports.EvaluationServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.EvaluationServiceRestInterceptor(), + ) + client = EvaluationServiceClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.EvaluationServiceRestInterceptor, "post_evaluate_instances" + ) as post, mock.patch.object( + transports.EvaluationServiceRestInterceptor, + "post_evaluate_instances_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.EvaluationServiceRestInterceptor, "pre_evaluate_instances" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = evaluation_service.EvaluateInstancesRequest.pb( + evaluation_service.EvaluateInstancesRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = evaluation_service.EvaluateInstancesResponse.to_json( + evaluation_service.EvaluateInstancesResponse() + ) + req.return_value.content = return_value - assert args[0] == request_msg + request = evaluation_service.EvaluateInstancesRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = evaluation_service.EvaluateInstancesResponse() + post_with_metadata.return_value = ( + evaluation_service.EvaluateInstancesResponse(), + metadata, + ) + client.evaluate_instances( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) -def test_transport_kind_rest(): - transport = EvaluationServiceClient.get_transport_class("rest")( - credentials=ga_credentials.AnonymousCredentials() - ) - assert transport.kind == "rest" + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() -def test_evaluate_instances_rest_bad_request( - request_type=evaluation_service.EvaluateInstancesRequest, +def test_evaluate_dataset_rest_bad_request( + request_type=evaluation_service.EvaluateDatasetRequest, ): client = EvaluationServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" @@ -1690,17 +2304,17 @@ def test_evaluate_instances_rest_bad_request( response_value.request = mock.Mock() req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.evaluate_instances(request) + client.evaluate_dataset(request) @pytest.mark.parametrize( "request_type", [ - evaluation_service.EvaluateInstancesRequest, + evaluation_service.EvaluateDatasetRequest, dict, ], ) -def test_evaluate_instances_rest_call_success(request_type): +def test_evaluate_dataset_rest_call_success(request_type): client = EvaluationServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" ) @@ -1712,26 +2326,23 @@ def test_evaluate_instances_rest_call_success(request_type): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = evaluation_service.EvaluateInstancesResponse() + return_value = operations_pb2.Operation(name="operations/spam") # Wrap the value into a proper Response obj response_value = mock.Mock() response_value.status_code = 200 - - # Convert return value to protobuf type - return_value = evaluation_service.EvaluateInstancesResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value.content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.evaluate_instances(request) + response = client.evaluate_dataset(request) # Establish that the response is the type that we expect. - assert isinstance(response, evaluation_service.EvaluateInstancesResponse) + json_return_value = json_format.MessageToJson(return_value) @pytest.mark.parametrize("null_interceptor", [True, False]) -def test_evaluate_instances_rest_interceptors(null_interceptor): +def test_evaluate_dataset_rest_interceptors(null_interceptor): transport = transports.EvaluationServiceRestTransport( credentials=ga_credentials.AnonymousCredentials(), interceptor=None @@ -1745,14 +2356,20 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): ) as req, mock.patch.object( path_template, "transcode" ) as transcode, mock.patch.object( - transports.EvaluationServiceRestInterceptor, "post_evaluate_instances" + operation.Operation, "_set_result_from_operation" + ), mock.patch.object( + transports.EvaluationServiceRestInterceptor, "post_evaluate_dataset" ) as post, mock.patch.object( - transports.EvaluationServiceRestInterceptor, "pre_evaluate_instances" + transports.EvaluationServiceRestInterceptor, + "post_evaluate_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.EvaluationServiceRestInterceptor, "pre_evaluate_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() - pb_message = evaluation_service.EvaluateInstancesRequest.pb( - evaluation_service.EvaluateInstancesRequest() + post_with_metadata.assert_not_called() + pb_message = evaluation_service.EvaluateDatasetRequest.pb( + evaluation_service.EvaluateDatasetRequest() ) transcode.return_value = { "method": "post", @@ -1764,20 +2381,19 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): req.return_value = mock.Mock() req.return_value.status_code = 200 req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - return_value = evaluation_service.EvaluateInstancesResponse.to_json( - evaluation_service.EvaluateInstancesResponse() - ) + return_value = json_format.MessageToJson(operations_pb2.Operation()) req.return_value.content = return_value - request = evaluation_service.EvaluateInstancesRequest() + request = evaluation_service.EvaluateDatasetRequest() metadata = [ ("key", "val"), ("cephalopod", "squid"), ] pre.return_value = request, metadata - post.return_value = evaluation_service.EvaluateInstancesResponse() + post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata - client.evaluate_instances( + client.evaluate_dataset( request, metadata=[ ("key", "val"), @@ -1787,6 +2403,7 @@ def test_evaluate_instances_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -2443,6 +3060,43 @@ def test_evaluate_instances_empty_call_rest(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_evaluate_dataset_empty_call_rest(): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + client.evaluate_dataset(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = evaluation_service.EvaluateDatasetRequest() + + assert args[0] == request_msg + + +def test_evaluation_service_rest_lro_client(): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + transport = client.transport + + # Ensure that we have an api-core operations client. + assert isinstance( + transport.operations_client, + operations_v1.AbstractOperationsClient, + ) + + # Ensure that subsequent calls to the property send the exact same object. + assert transport.operations_client is transport.operations_client + + def test_transport_kind_rest_asyncio(): if not HAS_ASYNC_REST_EXTRA: pytest.skip( @@ -2549,10 +3203,14 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncEvaluationServiceRestInterceptor, "post_evaluate_instances" ) as post, mock.patch.object( + transports.AsyncEvaluationServiceRestInterceptor, + "post_evaluate_instances_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncEvaluationServiceRestInterceptor, "pre_evaluate_instances" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = evaluation_service.EvaluateInstancesRequest.pb( evaluation_service.EvaluateInstancesRequest() ) @@ -2578,6 +3236,10 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = evaluation_service.EvaluateInstancesResponse() + post_with_metadata.return_value = ( + evaluation_service.EvaluateInstancesResponse(), + metadata, + ) await client.evaluate_instances( request, @@ -2589,6 +3251,147 @@ async def test_evaluate_instances_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() + + +@pytest.mark.asyncio +async def test_evaluate_dataset_rest_asyncio_bad_request( + request_type=evaluation_service.EvaluateDatasetRequest, +): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + # send a request that will satisfy transcoding + request_init = {"location": "projects/sample1/locations/sample2"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(AsyncAuthorizedSession, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.read = mock.AsyncMock(return_value=b"{}") + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + await client.evaluate_dataset(request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_type", + [ + evaluation_service.EvaluateDatasetRequest, + dict, + ], +) +async def test_evaluate_dataset_rest_asyncio_call_success(request_type): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + + # send a request that will satisfy transcoding + request_init = {"location": "projects/sample1/locations/sample2"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + response_value.read = mock.AsyncMock( + return_value=json_return_value.encode("UTF-8") + ) + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = await client.evaluate_dataset(request) + + # Establish that the response is the type that we expect. + json_return_value = json_format.MessageToJson(return_value) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("null_interceptor", [True, False]) +async def test_evaluate_dataset_rest_asyncio_interceptors(null_interceptor): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + transport = transports.AsyncEvaluationServiceRestTransport( + credentials=async_anonymous_credentials(), + interceptor=None + if null_interceptor + else transports.AsyncEvaluationServiceRestInterceptor(), + ) + client = EvaluationServiceAsyncClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + operation.Operation, "_set_result_from_operation" + ), mock.patch.object( + transports.AsyncEvaluationServiceRestInterceptor, "post_evaluate_dataset" + ) as post, mock.patch.object( + transports.AsyncEvaluationServiceRestInterceptor, + "post_evaluate_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.AsyncEvaluationServiceRestInterceptor, "pre_evaluate_dataset" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = evaluation_service.EvaluateDatasetRequest.pb( + evaluation_service.EvaluateDatasetRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = json_format.MessageToJson(operations_pb2.Operation()) + req.return_value.read = mock.AsyncMock(return_value=return_value) + + request = evaluation_service.EvaluateDatasetRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata + + await client.evaluate_dataset( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3366,6 +4169,52 @@ async def test_evaluate_instances_empty_call_rest_asyncio(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_evaluate_dataset_empty_call_rest_asyncio(): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="rest_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object(type(client.transport.evaluate_dataset), "__call__") as call: + await client.evaluate_dataset(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = evaluation_service.EvaluateDatasetRequest() + + assert args[0] == request_msg + + +def test_evaluation_service_rest_asyncio_lro_client(): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = EvaluationServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="rest_asyncio", + ) + transport = client.transport + + # Ensure that we have an api-core operations client. + assert isinstance( + transport.operations_client, + operations_v1.AsyncOperationsRestClient, + ) + + # Ensure that subsequent calls to the property send the exact same object. + assert transport.operations_client is transport.operations_client + + def test_unsupported_parameter_rest_asyncio(): if not HAS_ASYNC_REST_EXTRA: pytest.skip( @@ -3414,6 +4263,7 @@ def test_evaluation_service_base_transport(): # raise NotImplementedError. methods = ( "evaluate_instances", + "evaluate_dataset", "set_iam_policy", "get_iam_policy", "test_iam_permissions", @@ -3432,6 +4282,11 @@ def test_evaluation_service_base_transport(): with pytest.raises(NotImplementedError): transport.close() + # Additionally, the LRO client (a property) should + # also raise NotImplementedError + with pytest.raises(NotImplementedError): + transport.operations_client + # Catch all for all remaining methods and properties remainder = [ "kind", @@ -3685,6 +4540,9 @@ def test_evaluation_service_client_transport_session_collision(transport_name): session1 = client1.transport.evaluate_instances._session session2 = client2.transport.evaluate_instances._session assert session1 != session2 + session1 = client1.transport.evaluate_dataset._session + session2 = client2.transport.evaluate_dataset._session + assert session1 != session2 def test_evaluation_service_grpc_transport_channel(): @@ -3813,6 +4671,40 @@ def test_evaluation_service_transport_channel_mtls_with_adc(transport_class): assert transport.grpc_channel == mock_grpc_channel +def test_evaluation_service_grpc_lro_client(): + client = EvaluationServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + transport = client.transport + + # Ensure that we have a api-core operations client. + assert isinstance( + transport.operations_client, + operations_v1.OperationsClient, + ) + + # Ensure that subsequent calls to the property send the exact same object. + assert transport.operations_client is transport.operations_client + + +def test_evaluation_service_grpc_lro_async_client(): + client = EvaluationServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc_asyncio", + ) + transport = client.transport + + # Ensure that we have a api-core operations client. + assert isinstance( + transport.operations_client, + operations_v1.OperationsAsyncClient, + ) + + # Ensure that subsequent calls to the property send the exact same object. + assert transport.operations_client is transport.operations_client + + def test_common_billing_account_path(): billing_account = "squid" expected = "billingAccounts/{billing_account}".format( diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_extension_execution_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_extension_execution_service.py index 116d17025e..7091ccf4df 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_extension_execution_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_extension_execution_service.py @@ -86,6 +86,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -364,6 +372,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ExtensionExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ExtensionExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2554,10 +2605,14 @@ def test_execute_extension_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ExtensionExecutionServiceRestInterceptor, "post_execute_extension" ) as post, mock.patch.object( + transports.ExtensionExecutionServiceRestInterceptor, + "post_execute_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionExecutionServiceRestInterceptor, "pre_execute_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_execution_service.ExecuteExtensionRequest.pb( extension_execution_service.ExecuteExtensionRequest() ) @@ -2583,6 +2638,10 @@ def test_execute_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_execution_service.ExecuteExtensionResponse() + post_with_metadata.return_value = ( + extension_execution_service.ExecuteExtensionResponse(), + metadata, + ) client.execute_extension( request, @@ -2594,6 +2653,7 @@ def test_execute_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_extension_rest_bad_request( @@ -2680,10 +2740,14 @@ def test_query_extension_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ExtensionExecutionServiceRestInterceptor, "post_query_extension" ) as post, mock.patch.object( + transports.ExtensionExecutionServiceRestInterceptor, + "post_query_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionExecutionServiceRestInterceptor, "pre_query_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_execution_service.QueryExtensionRequest.pb( extension_execution_service.QueryExtensionRequest() ) @@ -2709,6 +2773,10 @@ def test_query_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_execution_service.QueryExtensionResponse() + post_with_metadata.return_value = ( + extension_execution_service.QueryExtensionResponse(), + metadata, + ) client.query_extension( request, @@ -2720,6 +2788,7 @@ def test_query_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3508,11 +3577,15 @@ async def test_execute_extension_rest_asyncio_interceptors(null_interceptor): transports.AsyncExtensionExecutionServiceRestInterceptor, "post_execute_extension", ) as post, mock.patch.object( + transports.AsyncExtensionExecutionServiceRestInterceptor, + "post_execute_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionExecutionServiceRestInterceptor, "pre_execute_extension", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_execution_service.ExecuteExtensionRequest.pb( extension_execution_service.ExecuteExtensionRequest() ) @@ -3538,6 +3611,10 @@ async def test_execute_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_execution_service.ExecuteExtensionResponse() + post_with_metadata.return_value = ( + extension_execution_service.ExecuteExtensionResponse(), + metadata, + ) await client.execute_extension( request, @@ -3549,6 +3626,7 @@ async def test_execute_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3651,10 +3729,14 @@ async def test_query_extension_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncExtensionExecutionServiceRestInterceptor, "post_query_extension" ) as post, mock.patch.object( + transports.AsyncExtensionExecutionServiceRestInterceptor, + "post_query_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionExecutionServiceRestInterceptor, "pre_query_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_execution_service.QueryExtensionRequest.pb( extension_execution_service.QueryExtensionRequest() ) @@ -3680,6 +3762,10 @@ async def test_query_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_execution_service.QueryExtensionResponse() + post_with_metadata.return_value = ( + extension_execution_service.QueryExtensionResponse(), + metadata, + ) await client.query_extension( request, @@ -3691,6 +3777,7 @@ async def test_query_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_extension_registry_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_extension_registry_service.py index 80c54e5bac..fe3299674b 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_extension_registry_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_extension_registry_service.py @@ -95,6 +95,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -369,6 +377,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ExtensionRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ExtensionRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4668,10 +4719,14 @@ def test_import_extension_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "post_import_extension" ) as post, mock.patch.object( + transports.ExtensionRegistryServiceRestInterceptor, + "post_import_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "pre_import_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.ImportExtensionRequest.pb( extension_registry_service.ImportExtensionRequest() ) @@ -4695,6 +4750,7 @@ def test_import_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_extension( request, @@ -4706,6 +4762,7 @@ def test_import_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_extension_rest_bad_request( @@ -4796,10 +4853,14 @@ def test_get_extension_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "post_get_extension" ) as post, mock.patch.object( + transports.ExtensionRegistryServiceRestInterceptor, + "post_get_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "pre_get_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.GetExtensionRequest.pb( extension_registry_service.GetExtensionRequest() ) @@ -4823,6 +4884,7 @@ def test_get_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension.Extension() + post_with_metadata.return_value = extension.Extension(), metadata client.get_extension( request, @@ -4834,6 +4896,7 @@ def test_get_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_extensions_rest_bad_request( @@ -4920,10 +4983,14 @@ def test_list_extensions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "post_list_extensions" ) as post, mock.patch.object( + transports.ExtensionRegistryServiceRestInterceptor, + "post_list_extensions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "pre_list_extensions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.ListExtensionsRequest.pb( extension_registry_service.ListExtensionsRequest() ) @@ -4949,6 +5016,10 @@ def test_list_extensions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_registry_service.ListExtensionsResponse() + post_with_metadata.return_value = ( + extension_registry_service.ListExtensionsResponse(), + metadata, + ) client.list_extensions( request, @@ -4960,6 +5031,7 @@ def test_list_extensions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_extension_rest_bad_request( @@ -5233,10 +5305,14 @@ def test_update_extension_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "post_update_extension" ) as post, mock.patch.object( + transports.ExtensionRegistryServiceRestInterceptor, + "post_update_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "pre_update_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.UpdateExtensionRequest.pb( extension_registry_service.UpdateExtensionRequest() ) @@ -5260,6 +5336,7 @@ def test_update_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_extension.Extension() + post_with_metadata.return_value = gca_extension.Extension(), metadata client.update_extension( request, @@ -5271,6 +5348,7 @@ def test_update_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_extension_rest_bad_request( @@ -5351,10 +5429,14 @@ def test_delete_extension_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "post_delete_extension" ) as post, mock.patch.object( + transports.ExtensionRegistryServiceRestInterceptor, + "post_delete_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ExtensionRegistryServiceRestInterceptor, "pre_delete_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.DeleteExtensionRequest.pb( extension_registry_service.DeleteExtensionRequest() ) @@ -5378,6 +5460,7 @@ def test_delete_extension_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_extension( request, @@ -5389,6 +5472,7 @@ def test_delete_extension_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6424,10 +6508,14 @@ async def test_import_extension_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "post_import_extension" ) as post, mock.patch.object( + transports.AsyncExtensionRegistryServiceRestInterceptor, + "post_import_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "pre_import_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.ImportExtensionRequest.pb( extension_registry_service.ImportExtensionRequest() ) @@ -6451,6 +6539,7 @@ async def test_import_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_extension( request, @@ -6462,6 +6551,7 @@ async def test_import_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6568,10 +6658,14 @@ async def test_get_extension_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "post_get_extension" ) as post, mock.patch.object( + transports.AsyncExtensionRegistryServiceRestInterceptor, + "post_get_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "pre_get_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.GetExtensionRequest.pb( extension_registry_service.GetExtensionRequest() ) @@ -6595,6 +6689,7 @@ async def test_get_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension.Extension() + post_with_metadata.return_value = extension.Extension(), metadata await client.get_extension( request, @@ -6606,6 +6701,7 @@ async def test_get_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6708,10 +6804,14 @@ async def test_list_extensions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "post_list_extensions" ) as post, mock.patch.object( + transports.AsyncExtensionRegistryServiceRestInterceptor, + "post_list_extensions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "pre_list_extensions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.ListExtensionsRequest.pb( extension_registry_service.ListExtensionsRequest() ) @@ -6737,6 +6837,10 @@ async def test_list_extensions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = extension_registry_service.ListExtensionsResponse() + post_with_metadata.return_value = ( + extension_registry_service.ListExtensionsResponse(), + metadata, + ) await client.list_extensions( request, @@ -6748,6 +6852,7 @@ async def test_list_extensions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7037,10 +7142,14 @@ async def test_update_extension_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "post_update_extension" ) as post, mock.patch.object( + transports.AsyncExtensionRegistryServiceRestInterceptor, + "post_update_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "pre_update_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.UpdateExtensionRequest.pb( extension_registry_service.UpdateExtensionRequest() ) @@ -7064,6 +7173,7 @@ async def test_update_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_extension.Extension() + post_with_metadata.return_value = gca_extension.Extension(), metadata await client.update_extension( request, @@ -7075,6 +7185,7 @@ async def test_update_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7171,10 +7282,14 @@ async def test_delete_extension_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "post_delete_extension" ) as post, mock.patch.object( + transports.AsyncExtensionRegistryServiceRestInterceptor, + "post_delete_extension_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncExtensionRegistryServiceRestInterceptor, "pre_delete_extension" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = extension_registry_service.DeleteExtensionRequest.pb( extension_registry_service.DeleteExtensionRequest() ) @@ -7198,6 +7313,7 @@ async def test_delete_extension_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_extension( request, @@ -7209,6 +7325,7 @@ async def test_delete_extension_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_admin_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_admin_service.py index ef695b7765..955511e0dd 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_admin_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_admin_service.py @@ -104,6 +104,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -393,6 +401,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureOnlineStoreAdminServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureOnlineStoreAdminServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -10353,11 +10404,15 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest() @@ -10383,6 +10438,7 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_online_store( request, @@ -10394,6 +10450,7 @@ def test_create_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_online_store_rest_bad_request( @@ -10491,11 +10548,15 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.GetFeatureOnlineStoreRequest() ) @@ -10521,6 +10582,10 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store.FeatureOnlineStore() + post_with_metadata.return_value = ( + feature_online_store.FeatureOnlineStore(), + metadata, + ) client.get_feature_online_store( request, @@ -10532,6 +10597,7 @@ def test_get_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_online_stores_rest_bad_request( @@ -10623,11 +10689,15 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_online_stores", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_online_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_online_stores", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.ListFeatureOnlineStoresRequest.pb( feature_online_store_admin_service.ListFeatureOnlineStoresRequest() @@ -10659,6 +10729,10 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureOnlineStoresResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureOnlineStoresResponse(), + metadata, + ) client.list_feature_online_stores( request, @@ -10670,6 +10744,7 @@ def test_list_feature_online_stores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_online_store_rest_bad_request( @@ -10865,11 +10940,15 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest() @@ -10895,6 +10974,7 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_online_store( request, @@ -10906,6 +10986,7 @@ def test_update_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_online_store_rest_bad_request( @@ -10991,11 +11072,15 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_online_store", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest() @@ -11021,6 +11106,7 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_online_store( request, @@ -11032,6 +11118,7 @@ def test_delete_feature_online_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_view_rest_bad_request( @@ -11241,11 +11328,15 @@ def test_create_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.CreateFeatureViewRequest.pb( feature_online_store_admin_service.CreateFeatureViewRequest() ) @@ -11269,6 +11360,7 @@ def test_create_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_view( request, @@ -11280,6 +11372,7 @@ def test_create_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_view_rest_bad_request( @@ -11382,10 +11475,14 @@ def test_get_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewRequest.pb( feature_online_store_admin_service.GetFeatureViewRequest() ) @@ -11409,6 +11506,7 @@ def test_get_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view.FeatureView() + post_with_metadata.return_value = feature_view.FeatureView(), metadata client.get_feature_view( request, @@ -11420,6 +11518,7 @@ def test_get_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_views_rest_bad_request( @@ -11511,11 +11610,15 @@ def test_list_feature_views_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_views", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_views_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_views", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewsRequest.pb( feature_online_store_admin_service.ListFeatureViewsRequest() ) @@ -11545,6 +11648,10 @@ def test_list_feature_views_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewsResponse(), + metadata, + ) client.list_feature_views( request, @@ -11556,6 +11663,7 @@ def test_list_feature_views_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_view_rest_bad_request( @@ -11769,11 +11877,15 @@ def test_update_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.UpdateFeatureViewRequest.pb( feature_online_store_admin_service.UpdateFeatureViewRequest() ) @@ -11797,6 +11909,7 @@ def test_update_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_view( request, @@ -11808,6 +11921,7 @@ def test_update_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_view_rest_bad_request( @@ -11893,11 +12007,15 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.DeleteFeatureViewRequest.pb( feature_online_store_admin_service.DeleteFeatureViewRequest() ) @@ -11921,6 +12039,7 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_view( request, @@ -11932,6 +12051,7 @@ def test_delete_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_sync_feature_view_rest_bad_request( @@ -12025,11 +12145,15 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_sync_feature_view", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_sync_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_sync_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.SyncFeatureViewRequest.pb( feature_online_store_admin_service.SyncFeatureViewRequest() ) @@ -12057,6 +12181,10 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_admin_service.SyncFeatureViewResponse() + post_with_metadata.return_value = ( + feature_online_store_admin_service.SyncFeatureViewResponse(), + metadata, + ) client.sync_feature_view( request, @@ -12068,6 +12196,7 @@ def test_sync_feature_view_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_view_sync_rest_bad_request( @@ -12161,11 +12290,15 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view_sync", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_sync_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view_sync", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewSyncRequest.pb( feature_online_store_admin_service.GetFeatureViewSyncRequest() ) @@ -12191,6 +12324,7 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view_sync.FeatureViewSync() + post_with_metadata.return_value = feature_view_sync.FeatureViewSync(), metadata client.get_feature_view_sync( request, @@ -12202,6 +12336,7 @@ def test_get_feature_view_sync_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_view_syncs_rest_bad_request( @@ -12295,11 +12430,15 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_view_syncs", ) as post, mock.patch.object( + transports.FeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_view_syncs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_view_syncs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewSyncsRequest.pb( feature_online_store_admin_service.ListFeatureViewSyncsRequest() ) @@ -12329,6 +12468,10 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewSyncsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewSyncsResponse(), + metadata, + ) client.list_feature_view_syncs( request, @@ -12340,6 +12483,7 @@ def test_list_feature_view_syncs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -13495,11 +13639,15 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.CreateFeatureOnlineStoreRequest() @@ -13525,6 +13673,7 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_online_store( request, @@ -13536,6 +13685,7 @@ async def test_create_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13649,11 +13799,15 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.GetFeatureOnlineStoreRequest() ) @@ -13679,6 +13833,10 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = feature_online_store.FeatureOnlineStore() + post_with_metadata.return_value = ( + feature_online_store.FeatureOnlineStore(), + metadata, + ) await client.get_feature_online_store( request, @@ -13690,6 +13848,7 @@ async def test_get_feature_online_store_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13797,11 +13956,15 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_online_stores", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_online_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_online_stores", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.ListFeatureOnlineStoresRequest.pb( feature_online_store_admin_service.ListFeatureOnlineStoresRequest() @@ -13833,6 +13996,10 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce post.return_value = ( feature_online_store_admin_service.ListFeatureOnlineStoresResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureOnlineStoresResponse(), + metadata, + ) await client.list_feature_online_stores( request, @@ -13844,6 +14011,7 @@ async def test_list_feature_online_stores_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14055,11 +14223,15 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.UpdateFeatureOnlineStoreRequest() @@ -14085,6 +14257,7 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_online_store( request, @@ -14096,6 +14269,7 @@ async def test_update_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14197,11 +14371,15 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_online_store", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_online_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_online_store", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest.pb( feature_online_store_admin_service.DeleteFeatureOnlineStoreRequest() @@ -14227,6 +14405,7 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_online_store( request, @@ -14238,6 +14417,7 @@ async def test_delete_feature_online_store_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14463,11 +14643,15 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_create_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_create_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_create_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.CreateFeatureViewRequest.pb( feature_online_store_admin_service.CreateFeatureViewRequest() ) @@ -14491,6 +14675,7 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_view( request, @@ -14502,6 +14687,7 @@ async def test_create_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14620,11 +14806,15 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewRequest.pb( feature_online_store_admin_service.GetFeatureViewRequest() ) @@ -14648,6 +14838,7 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_view.FeatureView() + post_with_metadata.return_value = feature_view.FeatureView(), metadata await client.get_feature_view( request, @@ -14659,6 +14850,7 @@ async def test_get_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14766,11 +14958,15 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_views", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_views_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_views", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewsRequest.pb( feature_online_store_admin_service.ListFeatureViewsRequest() ) @@ -14800,6 +14996,10 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): post.return_value = ( feature_online_store_admin_service.ListFeatureViewsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewsResponse(), + metadata, + ) await client.list_feature_views( request, @@ -14811,6 +15011,7 @@ async def test_list_feature_views_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15040,11 +15241,15 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_update_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_update_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_update_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.UpdateFeatureViewRequest.pb( feature_online_store_admin_service.UpdateFeatureViewRequest() ) @@ -15068,6 +15273,7 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_view( request, @@ -15079,6 +15285,7 @@ async def test_update_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15180,11 +15387,15 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_delete_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_delete_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_delete_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.DeleteFeatureViewRequest.pb( feature_online_store_admin_service.DeleteFeatureViewRequest() ) @@ -15208,6 +15419,7 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_view( request, @@ -15219,6 +15431,7 @@ async def test_delete_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15328,11 +15541,15 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_sync_feature_view", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_sync_feature_view_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_sync_feature_view", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.SyncFeatureViewRequest.pb( feature_online_store_admin_service.SyncFeatureViewRequest() ) @@ -15360,6 +15577,10 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_admin_service.SyncFeatureViewResponse() + post_with_metadata.return_value = ( + feature_online_store_admin_service.SyncFeatureViewResponse(), + metadata, + ) await client.sync_feature_view( request, @@ -15371,6 +15592,7 @@ async def test_sync_feature_view_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15480,11 +15702,15 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_get_feature_view_sync", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_get_feature_view_sync_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_get_feature_view_sync", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.GetFeatureViewSyncRequest.pb( feature_online_store_admin_service.GetFeatureViewSyncRequest() ) @@ -15510,6 +15736,7 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = feature_view_sync.FeatureViewSync() + post_with_metadata.return_value = feature_view_sync.FeatureViewSync(), metadata await client.get_feature_view_sync( request, @@ -15521,6 +15748,7 @@ async def test_get_feature_view_sync_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15630,11 +15858,15 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "post_list_feature_view_syncs", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, + "post_list_feature_view_syncs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreAdminServiceRestInterceptor, "pre_list_feature_view_syncs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_admin_service.ListFeatureViewSyncsRequest.pb( feature_online_store_admin_service.ListFeatureViewSyncsRequest() ) @@ -15664,6 +15896,10 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto post.return_value = ( feature_online_store_admin_service.ListFeatureViewSyncsResponse() ) + post_with_metadata.return_value = ( + feature_online_store_admin_service.ListFeatureViewSyncsResponse(), + metadata, + ) await client.list_feature_view_syncs( request, @@ -15675,6 +15911,7 @@ async def test_list_feature_view_syncs_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_service.py index e01328e072..be1824299e 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_feature_online_store_service.py @@ -83,6 +83,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -361,6 +369,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureOnlineStoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureOnlineStoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2580,10 +2631,14 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "post_fetch_feature_values" ) as post, mock.patch.object( + transports.FeatureOnlineStoreServiceRestInterceptor, + "post_fetch_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "pre_fetch_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.FetchFeatureValuesRequest.pb( feature_online_store_service.FetchFeatureValuesRequest() ) @@ -2609,6 +2664,10 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.FetchFeatureValuesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.FetchFeatureValuesResponse(), + metadata, + ) client.fetch_feature_values( request, @@ -2620,6 +2679,7 @@ def test_fetch_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_streaming_fetch_feature_values_rest_error(): @@ -2724,11 +2784,15 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): transports.FeatureOnlineStoreServiceRestInterceptor, "post_search_nearest_entities", ) as post, mock.patch.object( + transports.FeatureOnlineStoreServiceRestInterceptor, + "post_search_nearest_entities_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureOnlineStoreServiceRestInterceptor, "pre_search_nearest_entities", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.SearchNearestEntitiesRequest.pb( feature_online_store_service.SearchNearestEntitiesRequest() ) @@ -2756,6 +2820,10 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.SearchNearestEntitiesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.SearchNearestEntitiesResponse(), + metadata, + ) client.search_nearest_entities( request, @@ -2767,6 +2835,7 @@ def test_search_nearest_entities_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3558,11 +3627,15 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "post_fetch_feature_values", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreServiceRestInterceptor, + "post_fetch_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "pre_fetch_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.FetchFeatureValuesRequest.pb( feature_online_store_service.FetchFeatureValuesRequest() ) @@ -3588,6 +3661,10 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_online_store_service.FetchFeatureValuesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.FetchFeatureValuesResponse(), + metadata, + ) await client.fetch_feature_values( request, @@ -3599,6 +3676,7 @@ async def test_fetch_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3724,11 +3802,15 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "post_search_nearest_entities", ) as post, mock.patch.object( + transports.AsyncFeatureOnlineStoreServiceRestInterceptor, + "post_search_nearest_entities_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureOnlineStoreServiceRestInterceptor, "pre_search_nearest_entities", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_online_store_service.SearchNearestEntitiesRequest.pb( feature_online_store_service.SearchNearestEntitiesRequest() ) @@ -3756,6 +3838,10 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = feature_online_store_service.SearchNearestEntitiesResponse() + post_with_metadata.return_value = ( + feature_online_store_service.SearchNearestEntitiesResponse(), + metadata, + ) await client.search_nearest_entities( request, @@ -3767,6 +3853,7 @@ async def test_search_nearest_entities_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_feature_registry_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_feature_registry_service.py index 249d9d93fd..5efdfbb759 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_feature_registry_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_feature_registry_service.py @@ -107,6 +107,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -381,6 +389,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeatureRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeatureRegistryServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -13947,10 +13998,14 @@ def test_create_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_create_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureGroupRequest.pb( feature_registry_service.CreateFeatureGroupRequest() ) @@ -13974,6 +14029,7 @@ def test_create_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_group( request, @@ -13985,6 +14041,7 @@ def test_create_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_group_rest_bad_request( @@ -14080,10 +14137,14 @@ def test_get_feature_group_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureGroupRequest.pb( feature_registry_service.GetFeatureGroupRequest() ) @@ -14107,6 +14168,7 @@ def test_get_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_group.FeatureGroup() + post_with_metadata.return_value = feature_group.FeatureGroup(), metadata client.get_feature_group( request, @@ -14118,6 +14180,7 @@ def test_get_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_groups_rest_bad_request( @@ -14204,10 +14267,14 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_list_feature_groups" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_feature_groups_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_feature_groups" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureGroupsRequest.pb( feature_registry_service.ListFeatureGroupsRequest() ) @@ -14233,6 +14300,10 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureGroupsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureGroupsResponse(), + metadata, + ) client.list_feature_groups( request, @@ -14244,6 +14315,7 @@ def test_list_feature_groups_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_group_rest_bad_request( @@ -14421,10 +14493,14 @@ def test_update_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_update_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_update_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_update_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureGroupRequest.pb( feature_registry_service.UpdateFeatureGroupRequest() ) @@ -14448,6 +14524,7 @@ def test_update_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_group( request, @@ -14459,6 +14536,7 @@ def test_update_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_group_rest_bad_request( @@ -14539,10 +14617,14 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_delete_feature_group" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_delete_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_delete_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureGroupRequest.pb( feature_registry_service.DeleteFeatureGroupRequest() ) @@ -14566,6 +14648,7 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_group( request, @@ -14577,6 +14660,7 @@ def test_delete_feature_group_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_rest_bad_request( @@ -14782,10 +14866,14 @@ def test_create_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -14809,6 +14897,7 @@ def test_create_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature( request, @@ -14820,6 +14909,7 @@ def test_create_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_features_rest_bad_request( @@ -14904,10 +14994,14 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -14931,6 +15025,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_create_features( request, @@ -14942,6 +15037,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_rest_bad_request( @@ -15042,10 +15138,14 @@ def test_get_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -15069,6 +15169,7 @@ def test_get_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata client.get_feature( request, @@ -15080,6 +15181,7 @@ def test_get_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_features_rest_bad_request( @@ -15168,10 +15270,14 @@ def test_list_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -15197,6 +15303,10 @@ def test_list_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) client.list_features( request, @@ -15208,6 +15318,7 @@ def test_list_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_rest_bad_request( @@ -15417,10 +15528,14 @@ def test_update_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -15444,6 +15559,7 @@ def test_update_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature( request, @@ -15455,6 +15571,7 @@ def test_update_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_rest_bad_request( @@ -15539,10 +15656,14 @@ def test_delete_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -15566,6 +15687,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature( request, @@ -15577,6 +15699,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_monitor_rest_bad_request( @@ -15744,10 +15867,14 @@ def test_create_feature_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_create_feature_monitor" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureMonitorRequest.pb( feature_registry_service.CreateFeatureMonitorRequest() ) @@ -15771,6 +15898,7 @@ def test_create_feature_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature_monitor( request, @@ -15782,6 +15910,7 @@ def test_create_feature_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_monitor_rest_bad_request( @@ -15874,10 +16003,14 @@ def test_get_feature_monitor_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature_monitor" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureMonitorRequest.pb( feature_registry_service.GetFeatureMonitorRequest() ) @@ -15903,6 +16036,7 @@ def test_get_feature_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_monitor.FeatureMonitor() + post_with_metadata.return_value = feature_monitor.FeatureMonitor(), metadata client.get_feature_monitor( request, @@ -15914,6 +16048,7 @@ def test_get_feature_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_monitors_rest_bad_request( @@ -16004,10 +16139,14 @@ def test_list_feature_monitors_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_list_feature_monitors" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_feature_monitors_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_feature_monitors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureMonitorsRequest.pb( feature_registry_service.ListFeatureMonitorsRequest() ) @@ -16033,6 +16172,10 @@ def test_list_feature_monitors_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureMonitorsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureMonitorsResponse(), + metadata, + ) client.list_feature_monitors( request, @@ -16044,6 +16187,7 @@ def test_list_feature_monitors_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_monitor_rest_bad_request( @@ -16215,10 +16359,14 @@ def test_update_feature_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_update_feature_monitor" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_update_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_update_feature_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureMonitorRequest.pb( feature_registry_service.UpdateFeatureMonitorRequest() ) @@ -16242,6 +16390,7 @@ def test_update_feature_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_feature_monitor( request, @@ -16253,6 +16402,7 @@ def test_update_feature_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_monitor_rest_bad_request( @@ -16337,10 +16487,14 @@ def test_delete_feature_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_delete_feature_monitor" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_delete_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_delete_feature_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureMonitorRequest.pb( feature_registry_service.DeleteFeatureMonitorRequest() ) @@ -16364,6 +16518,7 @@ def test_delete_feature_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_monitor( request, @@ -16375,6 +16530,7 @@ def test_delete_feature_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_monitor_job_rest_bad_request( @@ -16588,11 +16744,15 @@ def test_create_feature_monitor_job_rest_interceptors(null_interceptor): transports.FeatureRegistryServiceRestInterceptor, "post_create_feature_monitor_job", ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_create_feature_monitor_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_create_feature_monitor_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureMonitorJobRequest.pb( feature_registry_service.CreateFeatureMonitorJobRequest() ) @@ -16618,6 +16778,10 @@ def test_create_feature_monitor_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_feature_monitor_job.FeatureMonitorJob() + post_with_metadata.return_value = ( + gca_feature_monitor_job.FeatureMonitorJob(), + metadata, + ) client.create_feature_monitor_job( request, @@ -16629,6 +16793,7 @@ def test_create_feature_monitor_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_monitor_job_rest_bad_request( @@ -16726,10 +16891,14 @@ def test_get_feature_monitor_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "post_get_feature_monitor_job" ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_get_feature_monitor_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_get_feature_monitor_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureMonitorJobRequest.pb( feature_registry_service.GetFeatureMonitorJobRequest() ) @@ -16755,6 +16924,10 @@ def test_get_feature_monitor_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_monitor_job.FeatureMonitorJob() + post_with_metadata.return_value = ( + feature_monitor_job.FeatureMonitorJob(), + metadata, + ) client.get_feature_monitor_job( request, @@ -16766,6 +16939,7 @@ def test_get_feature_monitor_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_feature_monitor_jobs_rest_bad_request( @@ -16857,11 +17031,15 @@ def test_list_feature_monitor_jobs_rest_interceptors(null_interceptor): transports.FeatureRegistryServiceRestInterceptor, "post_list_feature_monitor_jobs", ) as post, mock.patch.object( + transports.FeatureRegistryServiceRestInterceptor, + "post_list_feature_monitor_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeatureRegistryServiceRestInterceptor, "pre_list_feature_monitor_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureMonitorJobsRequest.pb( feature_registry_service.ListFeatureMonitorJobsRequest() ) @@ -16887,6 +17065,10 @@ def test_list_feature_monitor_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureMonitorJobsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureMonitorJobsResponse(), + metadata, + ) client.list_feature_monitor_jobs( request, @@ -16898,6 +17080,7 @@ def test_list_feature_monitor_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -18152,11 +18335,15 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureGroupRequest.pb( feature_registry_service.CreateFeatureGroupRequest() ) @@ -18180,6 +18367,7 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_group( request, @@ -18191,6 +18379,7 @@ async def test_create_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18302,10 +18491,14 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature_group" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature_group" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureGroupRequest.pb( feature_registry_service.GetFeatureGroupRequest() ) @@ -18329,6 +18522,7 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_group.FeatureGroup() + post_with_metadata.return_value = feature_group.FeatureGroup(), metadata await client.get_feature_group( request, @@ -18340,6 +18534,7 @@ async def test_get_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18443,10 +18638,14 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_feature_groups", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_feature_groups_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_feature_groups" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureGroupsRequest.pb( feature_registry_service.ListFeatureGroupsRequest() ) @@ -18472,6 +18671,10 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureGroupsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureGroupsResponse(), + metadata, + ) await client.list_feature_groups( request, @@ -18483,6 +18686,7 @@ async def test_list_feature_groups_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18677,11 +18881,15 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_update_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_update_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_update_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureGroupRequest.pb( feature_registry_service.UpdateFeatureGroupRequest() ) @@ -18705,6 +18913,7 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_group( request, @@ -18716,6 +18925,7 @@ async def test_update_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18813,11 +19023,15 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_delete_feature_group", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_delete_feature_group_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_delete_feature_group", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureGroupRequest.pb( feature_registry_service.DeleteFeatureGroupRequest() ) @@ -18841,6 +19055,7 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_group( request, @@ -18852,6 +19067,7 @@ async def test_delete_feature_group_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19073,10 +19289,14 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -19100,6 +19320,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature( request, @@ -19111,6 +19332,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19212,11 +19434,15 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) transports.AsyncFeatureRegistryServiceRestInterceptor, "post_batch_create_features", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_batch_create_features", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -19240,6 +19466,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_create_features( request, @@ -19251,6 +19478,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19367,10 +19595,14 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -19394,6 +19626,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata await client.get_feature( request, @@ -19405,6 +19638,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19509,10 +19743,14 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -19538,6 +19776,10 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) await client.list_features( request, @@ -19549,6 +19791,7 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19774,10 +20017,14 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -19801,6 +20048,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature( request, @@ -19812,6 +20060,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19912,10 +20161,14 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -19939,6 +20192,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature( request, @@ -19950,6 +20204,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20134,11 +20389,15 @@ async def test_create_feature_monitor_rest_asyncio_interceptors(null_interceptor transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature_monitor", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureMonitorRequest.pb( feature_registry_service.CreateFeatureMonitorRequest() ) @@ -20162,6 +20421,7 @@ async def test_create_feature_monitor_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature_monitor( request, @@ -20173,6 +20433,7 @@ async def test_create_feature_monitor_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20282,10 +20543,14 @@ async def test_get_feature_monitor_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature_monitor", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureMonitorRequest.pb( feature_registry_service.GetFeatureMonitorRequest() ) @@ -20311,6 +20576,7 @@ async def test_get_feature_monitor_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature_monitor.FeatureMonitor() + post_with_metadata.return_value = feature_monitor.FeatureMonitor(), metadata await client.get_feature_monitor( request, @@ -20322,6 +20588,7 @@ async def test_get_feature_monitor_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20429,11 +20696,15 @@ async def test_list_feature_monitors_rest_asyncio_interceptors(null_interceptor) transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_feature_monitors", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_feature_monitors_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_feature_monitors", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureMonitorsRequest.pb( feature_registry_service.ListFeatureMonitorsRequest() ) @@ -20459,6 +20730,10 @@ async def test_list_feature_monitors_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureMonitorsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureMonitorsResponse(), + metadata, + ) await client.list_feature_monitors( request, @@ -20470,6 +20745,7 @@ async def test_list_feature_monitors_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20658,11 +20934,15 @@ async def test_update_feature_monitor_rest_asyncio_interceptors(null_interceptor transports.AsyncFeatureRegistryServiceRestInterceptor, "post_update_feature_monitor", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_update_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_update_feature_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.UpdateFeatureMonitorRequest.pb( feature_registry_service.UpdateFeatureMonitorRequest() ) @@ -20686,6 +20966,7 @@ async def test_update_feature_monitor_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_feature_monitor( request, @@ -20697,6 +20978,7 @@ async def test_update_feature_monitor_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20798,11 +21080,15 @@ async def test_delete_feature_monitor_rest_asyncio_interceptors(null_interceptor transports.AsyncFeatureRegistryServiceRestInterceptor, "post_delete_feature_monitor", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_delete_feature_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_delete_feature_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.DeleteFeatureMonitorRequest.pb( feature_registry_service.DeleteFeatureMonitorRequest() ) @@ -20826,6 +21112,7 @@ async def test_delete_feature_monitor_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_monitor( request, @@ -20837,6 +21124,7 @@ async def test_delete_feature_monitor_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21066,11 +21354,15 @@ async def test_create_feature_monitor_job_rest_asyncio_interceptors(null_interce transports.AsyncFeatureRegistryServiceRestInterceptor, "post_create_feature_monitor_job", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_create_feature_monitor_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_create_feature_monitor_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.CreateFeatureMonitorJobRequest.pb( feature_registry_service.CreateFeatureMonitorJobRequest() ) @@ -21096,6 +21388,10 @@ async def test_create_feature_monitor_job_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = gca_feature_monitor_job.FeatureMonitorJob() + post_with_metadata.return_value = ( + gca_feature_monitor_job.FeatureMonitorJob(), + metadata, + ) await client.create_feature_monitor_job( request, @@ -21107,6 +21403,7 @@ async def test_create_feature_monitor_job_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21221,11 +21518,15 @@ async def test_get_feature_monitor_job_rest_asyncio_interceptors(null_intercepto transports.AsyncFeatureRegistryServiceRestInterceptor, "post_get_feature_monitor_job", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_get_feature_monitor_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_get_feature_monitor_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.GetFeatureMonitorJobRequest.pb( feature_registry_service.GetFeatureMonitorJobRequest() ) @@ -21251,6 +21552,10 @@ async def test_get_feature_monitor_job_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = feature_monitor_job.FeatureMonitorJob() + post_with_metadata.return_value = ( + feature_monitor_job.FeatureMonitorJob(), + metadata, + ) await client.get_feature_monitor_job( request, @@ -21262,6 +21567,7 @@ async def test_get_feature_monitor_job_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21369,11 +21675,15 @@ async def test_list_feature_monitor_jobs_rest_asyncio_interceptors(null_intercep transports.AsyncFeatureRegistryServiceRestInterceptor, "post_list_feature_monitor_jobs", ) as post, mock.patch.object( + transports.AsyncFeatureRegistryServiceRestInterceptor, + "post_list_feature_monitor_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeatureRegistryServiceRestInterceptor, "pre_list_feature_monitor_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = feature_registry_service.ListFeatureMonitorJobsRequest.pb( feature_registry_service.ListFeatureMonitorJobsRequest() ) @@ -21399,6 +21709,10 @@ async def test_list_feature_monitor_jobs_rest_asyncio_interceptors(null_intercep ] pre.return_value = request, metadata post.return_value = feature_registry_service.ListFeatureMonitorJobsResponse() + post_with_metadata.return_value = ( + feature_registry_service.ListFeatureMonitorJobsResponse(), + metadata, + ) await client.list_feature_monitor_jobs( request, @@ -21410,6 +21724,7 @@ async def test_list_feature_monitor_jobs_rest_asyncio_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_online_serving_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_online_serving_service.py index 39666f8c54..9853676b45 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_online_serving_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_online_serving_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -377,6 +385,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeaturestoreOnlineServingServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeaturestoreOnlineServingServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3228,11 +3279,15 @@ def test_read_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_read_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.ReadFeatureValuesRequest.pb( featurestore_online_service.ReadFeatureValuesRequest() ) @@ -3258,6 +3313,10 @@ def test_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) client.read_feature_values( request, @@ -3269,6 +3328,7 @@ def test_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_streaming_read_feature_values_rest_bad_request( @@ -3361,11 +3421,15 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_streaming_read_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_streaming_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_streaming_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.StreamingReadFeatureValuesRequest.pb( featurestore_online_service.StreamingReadFeatureValuesRequest() ) @@ -3391,6 +3455,10 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) client.streaming_read_feature_values( request, @@ -3402,6 +3470,7 @@ def test_streaming_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_feature_values_rest_bad_request( @@ -3490,11 +3559,15 @@ def test_write_feature_values_rest_interceptors(null_interceptor): transports.FeaturestoreOnlineServingServiceRestInterceptor, "post_write_feature_values", ) as post, mock.patch.object( + transports.FeaturestoreOnlineServingServiceRestInterceptor, + "post_write_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreOnlineServingServiceRestInterceptor, "pre_write_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.WriteFeatureValuesRequest.pb( featurestore_online_service.WriteFeatureValuesRequest() ) @@ -3520,6 +3593,10 @@ def test_write_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.WriteFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.WriteFeatureValuesResponse(), + metadata, + ) client.write_feature_values( request, @@ -3531,6 +3608,7 @@ def test_write_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -4344,11 +4422,15 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.ReadFeatureValuesRequest.pb( featurestore_online_service.ReadFeatureValuesRequest() ) @@ -4374,6 +4456,10 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) await client.read_feature_values( request, @@ -4385,6 +4471,7 @@ async def test_read_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4493,11 +4580,15 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_streaming_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_streaming_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_streaming_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.StreamingReadFeatureValuesRequest.pb( featurestore_online_service.StreamingReadFeatureValuesRequest() ) @@ -4523,6 +4614,10 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = featurestore_online_service.ReadFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.ReadFeatureValuesResponse(), + metadata, + ) await client.streaming_read_feature_values( request, @@ -4534,6 +4629,7 @@ async def test_streaming_read_feature_values_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4638,11 +4734,15 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "post_write_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, + "post_write_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreOnlineServingServiceRestInterceptor, "pre_write_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_online_service.WriteFeatureValuesRequest.pb( featurestore_online_service.WriteFeatureValuesRequest() ) @@ -4668,6 +4768,10 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_online_service.WriteFeatureValuesResponse() + post_with_metadata.return_value = ( + featurestore_online_service.WriteFeatureValuesResponse(), + metadata, + ) await client.write_feature_values( request, @@ -4679,6 +4783,7 @@ async def test_write_feature_values_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_service.py index b3904bd589..045b0cced8 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_featurestore_service.py @@ -103,6 +103,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -375,6 +383,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = FeaturestoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = FeaturestoreServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -14955,10 +15006,14 @@ def test_create_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeaturestoreRequest.pb( featurestore_service.CreateFeaturestoreRequest() ) @@ -14982,6 +15037,7 @@ def test_create_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_featurestore( request, @@ -14993,6 +15049,7 @@ def test_create_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_featurestore_rest_bad_request( @@ -15087,10 +15144,14 @@ def test_get_featurestore_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_get_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeaturestoreRequest.pb( featurestore_service.GetFeaturestoreRequest() ) @@ -15114,6 +15175,7 @@ def test_get_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore.Featurestore() + post_with_metadata.return_value = featurestore.Featurestore(), metadata client.get_featurestore( request, @@ -15125,6 +15187,7 @@ def test_get_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_featurestores_rest_bad_request( @@ -15209,10 +15272,14 @@ def test_list_featurestores_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_featurestores" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_featurestores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_featurestores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturestoresRequest.pb( featurestore_service.ListFeaturestoresRequest() ) @@ -15238,6 +15305,10 @@ def test_list_featurestores_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturestoresResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturestoresResponse(), + metadata, + ) client.list_featurestores( request, @@ -15249,6 +15320,7 @@ def test_list_featurestores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_featurestore_rest_bad_request( @@ -15426,10 +15498,14 @@ def test_update_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeaturestoreRequest.pb( featurestore_service.UpdateFeaturestoreRequest() ) @@ -15453,6 +15529,7 @@ def test_update_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_featurestore( request, @@ -15464,6 +15541,7 @@ def test_update_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_featurestore_rest_bad_request( @@ -15544,10 +15622,14 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_featurestore" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeaturestoreRequest.pb( featurestore_service.DeleteFeaturestoreRequest() ) @@ -15571,6 +15653,7 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_featurestore( request, @@ -15582,6 +15665,7 @@ def test_delete_featurestore_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_entity_type_rest_bad_request( @@ -15755,10 +15839,14 @@ def test_create_entity_type_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateEntityTypeRequest.pb( featurestore_service.CreateEntityTypeRequest() ) @@ -15782,6 +15870,7 @@ def test_create_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_entity_type( request, @@ -15793,6 +15882,7 @@ def test_create_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_entity_type_rest_bad_request( @@ -15891,10 +15981,14 @@ def test_get_entity_type_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_get_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetEntityTypeRequest.pb( featurestore_service.GetEntityTypeRequest() ) @@ -15918,6 +16012,7 @@ def test_get_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = entity_type.EntityType() + post_with_metadata.return_value = entity_type.EntityType(), metadata client.get_entity_type( request, @@ -15929,6 +16024,7 @@ def test_get_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_entity_types_rest_bad_request( @@ -16017,10 +16113,14 @@ def test_list_entity_types_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_entity_types" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_entity_types_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_entity_types" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListEntityTypesRequest.pb( featurestore_service.ListEntityTypesRequest() ) @@ -16046,6 +16146,10 @@ def test_list_entity_types_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListEntityTypesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListEntityTypesResponse(), + metadata, + ) client.list_entity_types( request, @@ -16057,6 +16161,7 @@ def test_list_entity_types_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_entity_type_rest_bad_request( @@ -16248,10 +16353,14 @@ def test_update_entity_type_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateEntityTypeRequest.pb( featurestore_service.UpdateEntityTypeRequest() ) @@ -16275,6 +16384,7 @@ def test_update_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_entity_type.EntityType() + post_with_metadata.return_value = gca_entity_type.EntityType(), metadata client.update_entity_type( request, @@ -16286,6 +16396,7 @@ def test_update_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_entity_type_rest_bad_request( @@ -16370,10 +16481,14 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_entity_type" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteEntityTypeRequest.pb( featurestore_service.DeleteEntityTypeRequest() ) @@ -16397,6 +16512,7 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_entity_type( request, @@ -16408,6 +16524,7 @@ def test_delete_entity_type_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_feature_rest_bad_request( @@ -16613,10 +16730,14 @@ def test_create_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -16640,6 +16761,7 @@ def test_create_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_feature( request, @@ -16651,6 +16773,7 @@ def test_create_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_features_rest_bad_request( @@ -16735,10 +16858,14 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -16762,6 +16889,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_create_features( request, @@ -16773,6 +16901,7 @@ def test_batch_create_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_feature_rest_bad_request( @@ -16873,10 +17002,13 @@ def test_get_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, "post_get_feature_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -16900,6 +17032,7 @@ def test_get_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata client.get_feature( request, @@ -16911,6 +17044,7 @@ def test_get_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_features_rest_bad_request( @@ -16999,10 +17133,14 @@ def test_list_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -17028,6 +17166,10 @@ def test_list_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) client.list_features( request, @@ -17039,6 +17181,7 @@ def test_list_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_feature_rest_bad_request( @@ -17264,10 +17407,14 @@ def test_update_feature_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -17291,6 +17438,7 @@ def test_update_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_feature.Feature() + post_with_metadata.return_value = gca_feature.Feature(), metadata client.update_feature( request, @@ -17302,6 +17450,7 @@ def test_update_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_rest_bad_request( @@ -17386,10 +17535,14 @@ def test_delete_feature_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -17413,6 +17566,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature( request, @@ -17424,6 +17578,7 @@ def test_delete_feature_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_feature_values_rest_bad_request( @@ -17508,10 +17663,14 @@ def test_import_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_import_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_import_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_import_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ImportFeatureValuesRequest.pb( featurestore_service.ImportFeatureValuesRequest() ) @@ -17535,6 +17694,7 @@ def test_import_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_feature_values( request, @@ -17546,6 +17706,7 @@ def test_import_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_read_feature_values_rest_bad_request( @@ -17630,10 +17791,14 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_batch_read_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_batch_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_batch_read_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchReadFeatureValuesRequest.pb( featurestore_service.BatchReadFeatureValuesRequest() ) @@ -17657,6 +17822,7 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_read_feature_values( request, @@ -17668,6 +17834,7 @@ def test_batch_read_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_feature_values_rest_bad_request( @@ -17752,10 +17919,14 @@ def test_export_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_export_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_export_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_export_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ExportFeatureValuesRequest.pb( featurestore_service.ExportFeatureValuesRequest() ) @@ -17779,6 +17950,7 @@ def test_export_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_feature_values( request, @@ -17790,6 +17962,7 @@ def test_export_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_feature_values_rest_bad_request( @@ -17874,10 +18047,14 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): ), mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_delete_feature_values" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_delete_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_delete_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureValuesRequest.pb( featurestore_service.DeleteFeatureValuesRequest() ) @@ -17901,6 +18078,7 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_feature_values( request, @@ -17912,6 +18090,7 @@ def test_delete_feature_values_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_features_rest_bad_request( @@ -17996,10 +18175,14 @@ def test_search_features_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "post_search_features" ) as post, mock.patch.object( + transports.FeaturestoreServiceRestInterceptor, + "post_search_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.FeaturestoreServiceRestInterceptor, "pre_search_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.SearchFeaturesRequest.pb( featurestore_service.SearchFeaturesRequest() ) @@ -18025,6 +18208,10 @@ def test_search_features_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.SearchFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.SearchFeaturesResponse(), + metadata, + ) client.search_features( request, @@ -18036,6 +18223,7 @@ def test_search_features_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -19327,10 +19515,14 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeaturestoreRequest.pb( featurestore_service.CreateFeaturestoreRequest() ) @@ -19354,6 +19546,7 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_featurestore( request, @@ -19365,6 +19558,7 @@ async def test_create_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19475,10 +19669,14 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeaturestoreRequest.pb( featurestore_service.GetFeaturestoreRequest() ) @@ -19502,6 +19700,7 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore.Featurestore() + post_with_metadata.return_value = featurestore.Featurestore(), metadata await client.get_featurestore( request, @@ -19513,6 +19712,7 @@ async def test_get_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19613,10 +19813,14 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_featurestores" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_featurestores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_featurestores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturestoresRequest.pb( featurestore_service.ListFeaturestoresRequest() ) @@ -19642,6 +19846,10 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturestoresResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturestoresResponse(), + metadata, + ) await client.list_featurestores( request, @@ -19653,6 +19861,7 @@ async def test_list_featurestores_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19846,10 +20055,14 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeaturestoreRequest.pb( featurestore_service.UpdateFeaturestoreRequest() ) @@ -19873,6 +20086,7 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_featurestore( request, @@ -19884,6 +20098,7 @@ async def test_update_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19980,10 +20195,14 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_featurestore" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_featurestore_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_featurestore" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeaturestoreRequest.pb( featurestore_service.DeleteFeaturestoreRequest() ) @@ -20007,6 +20226,7 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_featurestore( request, @@ -20018,6 +20238,7 @@ async def test_delete_featurestore_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20207,10 +20428,14 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateEntityTypeRequest.pb( featurestore_service.CreateEntityTypeRequest() ) @@ -20234,6 +20459,7 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_entity_type( request, @@ -20245,6 +20471,7 @@ async def test_create_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20359,10 +20586,14 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetEntityTypeRequest.pb( featurestore_service.GetEntityTypeRequest() ) @@ -20386,6 +20617,7 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = entity_type.EntityType() + post_with_metadata.return_value = entity_type.EntityType(), metadata await client.get_entity_type( request, @@ -20397,6 +20629,7 @@ async def test_get_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20501,10 +20734,14 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_entity_types" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_entity_types_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_entity_types" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListEntityTypesRequest.pb( featurestore_service.ListEntityTypesRequest() ) @@ -20530,6 +20767,10 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListEntityTypesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListEntityTypesResponse(), + metadata, + ) await client.list_entity_types( request, @@ -20541,6 +20782,7 @@ async def test_list_entity_types_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20748,10 +20990,14 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateEntityTypeRequest.pb( featurestore_service.UpdateEntityTypeRequest() ) @@ -20775,6 +21021,7 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_entity_type.EntityType() + post_with_metadata.return_value = gca_entity_type.EntityType(), metadata await client.update_entity_type( request, @@ -20786,6 +21033,7 @@ async def test_update_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20886,10 +21134,14 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_entity_type" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_entity_type_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_entity_type" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteEntityTypeRequest.pb( featurestore_service.DeleteEntityTypeRequest() ) @@ -20913,6 +21165,7 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_entity_type( request, @@ -20924,6 +21177,7 @@ async def test_delete_entity_type_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21145,10 +21399,14 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_create_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_create_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_create_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.CreateFeatureRequest.pb( featurestore_service.CreateFeatureRequest() ) @@ -21172,6 +21430,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_feature( request, @@ -21183,6 +21442,7 @@ async def test_create_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21283,10 +21543,14 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_batch_create_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_batch_create_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_batch_create_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchCreateFeaturesRequest.pb( featurestore_service.BatchCreateFeaturesRequest() ) @@ -21310,6 +21574,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_create_features( request, @@ -21321,6 +21586,7 @@ async def test_batch_create_features_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21437,10 +21703,14 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_get_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_get_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_get_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.GetFeatureRequest.pb( featurestore_service.GetFeatureRequest() ) @@ -21464,6 +21734,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = feature.Feature() + post_with_metadata.return_value = feature.Feature(), metadata await client.get_feature( request, @@ -21475,6 +21746,7 @@ async def test_get_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21579,10 +21851,14 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_list_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_list_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_list_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ListFeaturesRequest.pb( featurestore_service.ListFeaturesRequest() ) @@ -21608,6 +21884,10 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.ListFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.ListFeaturesResponse(), + metadata, + ) await client.list_features( request, @@ -21619,6 +21899,7 @@ async def test_list_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21860,10 +22141,14 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_update_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_update_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_update_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.UpdateFeatureRequest.pb( featurestore_service.UpdateFeatureRequest() ) @@ -21887,6 +22172,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_feature.Feature() + post_with_metadata.return_value = gca_feature.Feature(), metadata await client.update_feature( request, @@ -21898,6 +22184,7 @@ async def test_update_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -21998,10 +22285,14 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_feature" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_feature_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_feature" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureRequest.pb( featurestore_service.DeleteFeatureRequest() ) @@ -22025,6 +22316,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature( request, @@ -22036,6 +22328,7 @@ async def test_delete_feature_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22136,10 +22429,14 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_import_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_import_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_import_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ImportFeatureValuesRequest.pb( featurestore_service.ImportFeatureValuesRequest() ) @@ -22163,6 +22460,7 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_feature_values( request, @@ -22174,6 +22472,7 @@ async def test_import_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22275,11 +22574,15 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep transports.AsyncFeaturestoreServiceRestInterceptor, "post_batch_read_feature_values", ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_batch_read_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_batch_read_feature_values", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.BatchReadFeatureValuesRequest.pb( featurestore_service.BatchReadFeatureValuesRequest() ) @@ -22303,6 +22606,7 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_read_feature_values( request, @@ -22314,6 +22618,7 @@ async def test_batch_read_feature_values_rest_asyncio_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22414,10 +22719,14 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_export_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_export_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_export_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.ExportFeatureValuesRequest.pb( featurestore_service.ExportFeatureValuesRequest() ) @@ -22441,6 +22750,7 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_feature_values( request, @@ -22452,6 +22762,7 @@ async def test_export_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22552,10 +22863,14 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_delete_feature_values" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_delete_feature_values_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_delete_feature_values" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.DeleteFeatureValuesRequest.pb( featurestore_service.DeleteFeatureValuesRequest() ) @@ -22579,6 +22894,7 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_feature_values( request, @@ -22590,6 +22906,7 @@ async def test_delete_feature_values_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -22690,10 +23007,14 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "post_search_features" ) as post, mock.patch.object( + transports.AsyncFeaturestoreServiceRestInterceptor, + "post_search_features_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncFeaturestoreServiceRestInterceptor, "pre_search_features" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = featurestore_service.SearchFeaturesRequest.pb( featurestore_service.SearchFeaturesRequest() ) @@ -22719,6 +23040,10 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = featurestore_service.SearchFeaturesResponse() + post_with_metadata.return_value = ( + featurestore_service.SearchFeaturesResponse(), + metadata, + ) await client.search_features( request, @@ -22730,6 +23055,7 @@ async def test_search_features_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_cache_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_cache_service.py index 851edc2866..0b5fc055ef 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_cache_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_cache_service.py @@ -90,6 +90,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -356,6 +364,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = GenAiCacheServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = GenAiCacheServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4796,10 +4847,14 @@ def test_create_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_create_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_create_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_create_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.CreateCachedContentRequest.pb( gen_ai_cache_service.CreateCachedContentRequest() ) @@ -4825,6 +4880,7 @@ def test_create_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata client.create_cached_content( request, @@ -4836,6 +4892,7 @@ def test_create_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_cached_content_rest_bad_request( @@ -4924,10 +4981,14 @@ def test_get_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_get_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_get_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_get_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.GetCachedContentRequest.pb( gen_ai_cache_service.GetCachedContentRequest() ) @@ -4953,6 +5014,7 @@ def test_get_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = cached_content.CachedContent() + post_with_metadata.return_value = cached_content.CachedContent(), metadata client.get_cached_content( request, @@ -4964,6 +5026,7 @@ def test_get_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_cached_content_rest_bad_request( @@ -5262,10 +5325,14 @@ def test_update_cached_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_update_cached_content" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_update_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_update_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.UpdateCachedContentRequest.pb( gen_ai_cache_service.UpdateCachedContentRequest() ) @@ -5291,6 +5358,7 @@ def test_update_cached_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata client.update_cached_content( request, @@ -5302,6 +5370,7 @@ def test_update_cached_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_cached_content_rest_bad_request( @@ -5495,10 +5564,14 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "post_list_cached_contents" ) as post, mock.patch.object( + transports.GenAiCacheServiceRestInterceptor, + "post_list_cached_contents_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiCacheServiceRestInterceptor, "pre_list_cached_contents" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.ListCachedContentsRequest.pb( gen_ai_cache_service.ListCachedContentsRequest() ) @@ -5524,6 +5597,10 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gen_ai_cache_service.ListCachedContentsResponse() + post_with_metadata.return_value = ( + gen_ai_cache_service.ListCachedContentsResponse(), + metadata, + ) client.list_cached_contents( request, @@ -5535,6 +5612,7 @@ def test_list_cached_contents_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6594,10 +6672,14 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_create_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_create_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_create_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.CreateCachedContentRequest.pb( gen_ai_cache_service.CreateCachedContentRequest() ) @@ -6623,6 +6705,7 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata await client.create_cached_content( request, @@ -6634,6 +6717,7 @@ async def test_create_cached_content_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6738,10 +6822,14 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_get_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_get_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_get_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.GetCachedContentRequest.pb( gen_ai_cache_service.GetCachedContentRequest() ) @@ -6767,6 +6855,7 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = cached_content.CachedContent() + post_with_metadata.return_value = cached_content.CachedContent(), metadata await client.get_cached_content( request, @@ -6778,6 +6867,7 @@ async def test_get_cached_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7092,10 +7182,14 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_update_cached_content" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_update_cached_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_update_cached_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.UpdateCachedContentRequest.pb( gen_ai_cache_service.UpdateCachedContentRequest() ) @@ -7121,6 +7215,7 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_cached_content.CachedContent() + post_with_metadata.return_value = gca_cached_content.CachedContent(), metadata await client.update_cached_content( request, @@ -7132,6 +7227,7 @@ async def test_update_cached_content_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7357,10 +7453,14 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "post_list_cached_contents" ) as post, mock.patch.object( + transports.AsyncGenAiCacheServiceRestInterceptor, + "post_list_cached_contents_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiCacheServiceRestInterceptor, "pre_list_cached_contents" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = gen_ai_cache_service.ListCachedContentsRequest.pb( gen_ai_cache_service.ListCachedContentsRequest() ) @@ -7386,6 +7486,10 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gen_ai_cache_service.ListCachedContentsResponse() + post_with_metadata.return_value = ( + gen_ai_cache_service.ListCachedContentsResponse(), + metadata, + ) await client.list_cached_contents( request, @@ -7397,6 +7501,7 @@ async def test_list_cached_contents_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_tuning_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_tuning_service.py index 3591b310fe..a3d5c0c54c 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_tuning_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_gen_ai_tuning_service.py @@ -96,6 +96,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -366,6 +374,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = GenAiTuningServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = GenAiTuningServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4755,10 +4806,14 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_create_tuning_job" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_create_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_create_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.CreateTuningJobRequest.pb( genai_tuning_service.CreateTuningJobRequest() ) @@ -4782,6 +4837,7 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tuning_job.TuningJob() + post_with_metadata.return_value = gca_tuning_job.TuningJob(), metadata client.create_tuning_job( request, @@ -4793,6 +4849,7 @@ def test_create_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tuning_job_rest_bad_request( @@ -4890,10 +4947,14 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_get_tuning_job" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_get_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_get_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.GetTuningJobRequest.pb( genai_tuning_service.GetTuningJobRequest() ) @@ -4917,6 +4978,7 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tuning_job.TuningJob() + post_with_metadata.return_value = tuning_job.TuningJob(), metadata client.get_tuning_job( request, @@ -4928,6 +4990,7 @@ def test_get_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tuning_jobs_rest_bad_request( @@ -5012,10 +5075,14 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_list_tuning_jobs" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_list_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_list_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.ListTuningJobsRequest.pb( genai_tuning_service.ListTuningJobsRequest() ) @@ -5041,6 +5108,10 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = genai_tuning_service.ListTuningJobsResponse() + post_with_metadata.return_value = ( + genai_tuning_service.ListTuningJobsResponse(), + metadata, + ) client.list_tuning_jobs( request, @@ -5052,6 +5123,7 @@ def test_list_tuning_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_tuning_job_rest_bad_request( @@ -5241,10 +5313,14 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "post_rebase_tuned_model" ) as post, mock.patch.object( + transports.GenAiTuningServiceRestInterceptor, + "post_rebase_tuned_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.GenAiTuningServiceRestInterceptor, "pre_rebase_tuned_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.RebaseTunedModelRequest.pb( genai_tuning_service.RebaseTunedModelRequest() ) @@ -5268,6 +5344,7 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.rebase_tuned_model( request, @@ -5279,6 +5356,7 @@ def test_rebase_tuned_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6363,10 +6441,14 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_create_tuning_job" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_create_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_create_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.CreateTuningJobRequest.pb( genai_tuning_service.CreateTuningJobRequest() ) @@ -6390,6 +6472,7 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tuning_job.TuningJob() + post_with_metadata.return_value = gca_tuning_job.TuningJob(), metadata await client.create_tuning_job( request, @@ -6401,6 +6484,7 @@ async def test_create_tuning_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6514,10 +6598,14 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_get_tuning_job" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_get_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_get_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.GetTuningJobRequest.pb( genai_tuning_service.GetTuningJobRequest() ) @@ -6541,6 +6629,7 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tuning_job.TuningJob() + post_with_metadata.return_value = tuning_job.TuningJob(), metadata await client.get_tuning_job( request, @@ -6552,6 +6641,7 @@ async def test_get_tuning_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6652,10 +6742,14 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_list_tuning_jobs" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_list_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_list_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.ListTuningJobsRequest.pb( genai_tuning_service.ListTuningJobsRequest() ) @@ -6681,6 +6775,10 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = genai_tuning_service.ListTuningJobsResponse() + post_with_metadata.return_value = ( + genai_tuning_service.ListTuningJobsResponse(), + metadata, + ) await client.list_tuning_jobs( request, @@ -6692,6 +6790,7 @@ async def test_list_tuning_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6913,10 +7012,14 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "post_rebase_tuned_model" ) as post, mock.patch.object( + transports.AsyncGenAiTuningServiceRestInterceptor, + "post_rebase_tuned_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncGenAiTuningServiceRestInterceptor, "pre_rebase_tuned_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = genai_tuning_service.RebaseTunedModelRequest.pb( genai_tuning_service.RebaseTunedModelRequest() ) @@ -6940,6 +7043,7 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.rebase_tuned_model( request, @@ -6951,6 +7055,7 @@ async def test_rebase_tuned_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_index_endpoint_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_index_endpoint_service.py index b6b7ba7f16..ca17671822 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_index_endpoint_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_index_endpoint_service.py @@ -95,6 +95,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -367,6 +375,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = IndexEndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = IndexEndpointServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -6585,10 +6636,14 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_create_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_create_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_create_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.CreateIndexEndpointRequest.pb( index_endpoint_service.CreateIndexEndpointRequest() ) @@ -6612,6 +6667,7 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_index_endpoint( request, @@ -6623,6 +6679,7 @@ def test_create_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_index_endpoint_rest_bad_request( @@ -6725,10 +6782,14 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_get_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_get_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_get_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.GetIndexEndpointRequest.pb( index_endpoint_service.GetIndexEndpointRequest() ) @@ -6754,6 +6815,7 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint.IndexEndpoint() + post_with_metadata.return_value = index_endpoint.IndexEndpoint(), metadata client.get_index_endpoint( request, @@ -6765,6 +6827,7 @@ def test_get_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_index_endpoints_rest_bad_request( @@ -6851,10 +6914,14 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_list_index_endpoints" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_list_index_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_list_index_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.ListIndexEndpointsRequest.pb( index_endpoint_service.ListIndexEndpointsRequest() ) @@ -6880,6 +6947,10 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint_service.ListIndexEndpointsResponse() + post_with_metadata.return_value = ( + index_endpoint_service.ListIndexEndpointsResponse(), + metadata, + ) client.list_index_endpoints( request, @@ -6891,6 +6962,7 @@ def test_list_index_endpoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_index_endpoint_rest_bad_request( @@ -7157,10 +7229,14 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_update_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_update_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_update_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UpdateIndexEndpointRequest.pb( index_endpoint_service.UpdateIndexEndpointRequest() ) @@ -7186,6 +7262,7 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_index_endpoint.IndexEndpoint() + post_with_metadata.return_value = gca_index_endpoint.IndexEndpoint(), metadata client.update_index_endpoint( request, @@ -7197,6 +7274,7 @@ def test_update_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_index_endpoint_rest_bad_request( @@ -7277,10 +7355,14 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_delete_index_endpoint" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_delete_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_delete_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeleteIndexEndpointRequest.pb( index_endpoint_service.DeleteIndexEndpointRequest() ) @@ -7304,6 +7386,7 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_index_endpoint( request, @@ -7315,6 +7398,7 @@ def test_delete_index_endpoint_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_deploy_index_rest_bad_request( @@ -7399,10 +7483,14 @@ def test_deploy_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_deploy_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_deploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_deploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeployIndexRequest.pb( index_endpoint_service.DeployIndexRequest() ) @@ -7426,6 +7514,7 @@ def test_deploy_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.deploy_index( request, @@ -7437,6 +7526,7 @@ def test_deploy_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_undeploy_index_rest_bad_request( @@ -7521,10 +7611,14 @@ def test_undeploy_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_undeploy_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_undeploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_undeploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UndeployIndexRequest.pb( index_endpoint_service.UndeployIndexRequest() ) @@ -7548,6 +7642,7 @@ def test_undeploy_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.undeploy_index( request, @@ -7559,6 +7654,7 @@ def test_undeploy_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_mutate_deployed_index_rest_bad_request( @@ -7766,10 +7862,14 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "post_mutate_deployed_index" ) as post, mock.patch.object( + transports.IndexEndpointServiceRestInterceptor, + "post_mutate_deployed_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.IndexEndpointServiceRestInterceptor, "pre_mutate_deployed_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.MutateDeployedIndexRequest.pb( index_endpoint_service.MutateDeployedIndexRequest() ) @@ -7793,6 +7893,7 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.mutate_deployed_index( request, @@ -7804,6 +7905,7 @@ def test_mutate_deployed_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -8889,10 +8991,14 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_create_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_create_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_create_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.CreateIndexEndpointRequest.pb( index_endpoint_service.CreateIndexEndpointRequest() ) @@ -8916,6 +9022,7 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_index_endpoint( request, @@ -8927,6 +9034,7 @@ async def test_create_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9045,10 +9153,14 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_get_index_endpoint" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_get_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_get_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.GetIndexEndpointRequest.pb( index_endpoint_service.GetIndexEndpointRequest() ) @@ -9074,6 +9186,7 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint.IndexEndpoint() + post_with_metadata.return_value = index_endpoint.IndexEndpoint(), metadata await client.get_index_endpoint( request, @@ -9085,6 +9198,7 @@ async def test_get_index_endpoint_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9187,10 +9301,14 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_list_index_endpoints" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_list_index_endpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_list_index_endpoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.ListIndexEndpointsRequest.pb( index_endpoint_service.ListIndexEndpointsRequest() ) @@ -9216,6 +9334,10 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_endpoint_service.ListIndexEndpointsResponse() + post_with_metadata.return_value = ( + index_endpoint_service.ListIndexEndpointsResponse(), + metadata, + ) await client.list_index_endpoints( request, @@ -9227,6 +9349,7 @@ async def test_list_index_endpoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9510,10 +9633,14 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_update_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_update_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_update_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UpdateIndexEndpointRequest.pb( index_endpoint_service.UpdateIndexEndpointRequest() ) @@ -9539,6 +9666,7 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = gca_index_endpoint.IndexEndpoint() + post_with_metadata.return_value = gca_index_endpoint.IndexEndpoint(), metadata await client.update_index_endpoint( request, @@ -9550,6 +9678,7 @@ async def test_update_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9647,10 +9776,14 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_delete_index_endpoint", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_delete_index_endpoint_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_delete_index_endpoint" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeleteIndexEndpointRequest.pb( index_endpoint_service.DeleteIndexEndpointRequest() ) @@ -9674,6 +9807,7 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_index_endpoint( request, @@ -9685,6 +9819,7 @@ async def test_delete_index_endpoint_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9785,10 +9920,14 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_deploy_index" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_deploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_deploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.DeployIndexRequest.pb( index_endpoint_service.DeployIndexRequest() ) @@ -9812,6 +9951,7 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.deploy_index( request, @@ -9823,6 +9963,7 @@ async def test_deploy_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9923,10 +10064,14 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "post_undeploy_index" ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_undeploy_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_undeploy_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.UndeployIndexRequest.pb( index_endpoint_service.UndeployIndexRequest() ) @@ -9950,6 +10095,7 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.undeploy_index( request, @@ -9961,6 +10107,7 @@ async def test_undeploy_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10185,10 +10332,14 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) transports.AsyncIndexEndpointServiceRestInterceptor, "post_mutate_deployed_index", ) as post, mock.patch.object( + transports.AsyncIndexEndpointServiceRestInterceptor, + "post_mutate_deployed_index_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexEndpointServiceRestInterceptor, "pre_mutate_deployed_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_endpoint_service.MutateDeployedIndexRequest.pb( index_endpoint_service.MutateDeployedIndexRequest() ) @@ -10212,6 +10363,7 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.mutate_deployed_index( request, @@ -10223,6 +10375,7 @@ async def test_mutate_deployed_index_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_index_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_index_service.py index a370e0df4f..6cdeb3dc3d 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_index_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_index_service.py @@ -91,6 +91,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -333,6 +341,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = IndexServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = IndexServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5293,10 +5344,13 @@ def test_create_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_create_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_create_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_create_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.CreateIndexRequest.pb( index_service.CreateIndexRequest() ) @@ -5320,6 +5374,7 @@ def test_create_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_index( request, @@ -5331,6 +5386,7 @@ def test_create_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_index_rest_bad_request(request_type=index_service.GetIndexRequest): @@ -5427,10 +5483,13 @@ def test_get_index_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_get_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_get_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_get_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.GetIndexRequest.pb(index_service.GetIndexRequest()) transcode.return_value = { "method": "post", @@ -5452,6 +5511,7 @@ def test_get_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index.Index() + post_with_metadata.return_value = index.Index(), metadata client.get_index( request, @@ -5463,6 +5523,7 @@ def test_get_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_indexes_rest_bad_request(request_type=index_service.ListIndexesRequest): @@ -5545,10 +5606,13 @@ def test_list_indexes_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_list_indexes" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_list_indexes_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_list_indexes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.ListIndexesRequest.pb( index_service.ListIndexesRequest() ) @@ -5574,6 +5638,7 @@ def test_list_indexes_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.ListIndexesResponse() + post_with_metadata.return_value = index_service.ListIndexesResponse(), metadata client.list_indexes( request, @@ -5585,6 +5650,7 @@ def test_list_indexes_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_index_rest_bad_request(request_type=index_service.UpdateIndexRequest): @@ -5768,10 +5834,13 @@ def test_update_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_update_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_update_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_update_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpdateIndexRequest.pb( index_service.UpdateIndexRequest() ) @@ -5795,6 +5864,7 @@ def test_update_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_index( request, @@ -5806,6 +5876,7 @@ def test_update_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_index_rest_bad_request(request_type=index_service.DeleteIndexRequest): @@ -5884,10 +5955,13 @@ def test_delete_index_rest_interceptors(null_interceptor): ), mock.patch.object( transports.IndexServiceRestInterceptor, "post_delete_index" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_delete_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_delete_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.DeleteIndexRequest.pb( index_service.DeleteIndexRequest() ) @@ -5911,6 +5985,7 @@ def test_delete_index_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_index( request, @@ -5922,6 +5997,7 @@ def test_delete_index_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upsert_datapoints_rest_bad_request( @@ -6003,10 +6079,13 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_upsert_datapoints" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_upsert_datapoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_upsert_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpsertDatapointsRequest.pb( index_service.UpsertDatapointsRequest() ) @@ -6032,6 +6111,10 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.UpsertDatapointsResponse() + post_with_metadata.return_value = ( + index_service.UpsertDatapointsResponse(), + metadata, + ) client.upsert_datapoints( request, @@ -6043,6 +6126,7 @@ def test_upsert_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_remove_datapoints_rest_bad_request( @@ -6124,10 +6208,13 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.IndexServiceRestInterceptor, "post_remove_datapoints" ) as post, mock.patch.object( + transports.IndexServiceRestInterceptor, "post_remove_datapoints_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.IndexServiceRestInterceptor, "pre_remove_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.RemoveDatapointsRequest.pb( index_service.RemoveDatapointsRequest() ) @@ -6153,6 +6240,10 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.RemoveDatapointsResponse() + post_with_metadata.return_value = ( + index_service.RemoveDatapointsResponse(), + metadata, + ) client.remove_datapoints( request, @@ -6164,6 +6255,7 @@ def test_remove_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7165,10 +7257,13 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_create_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_create_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_create_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.CreateIndexRequest.pb( index_service.CreateIndexRequest() ) @@ -7192,6 +7287,7 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_index( request, @@ -7203,6 +7299,7 @@ async def test_create_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7317,10 +7414,13 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_get_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_get_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_get_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.GetIndexRequest.pb(index_service.GetIndexRequest()) transcode.return_value = { "method": "post", @@ -7342,6 +7442,7 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index.Index() + post_with_metadata.return_value = index.Index(), metadata await client.get_index( request, @@ -7353,6 +7454,7 @@ async def test_get_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7453,10 +7555,13 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_list_indexes" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_list_indexes_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_list_indexes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.ListIndexesRequest.pb( index_service.ListIndexesRequest() ) @@ -7482,6 +7587,7 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.ListIndexesResponse() + post_with_metadata.return_value = index_service.ListIndexesResponse(), metadata await client.list_indexes( request, @@ -7493,6 +7599,7 @@ async def test_list_indexes_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7694,10 +7801,13 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_update_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_update_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_update_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpdateIndexRequest.pb( index_service.UpdateIndexRequest() ) @@ -7721,6 +7831,7 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_index( request, @@ -7732,6 +7843,7 @@ async def test_update_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7828,10 +7940,13 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_delete_index" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, "post_delete_index_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_delete_index" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.DeleteIndexRequest.pb( index_service.DeleteIndexRequest() ) @@ -7855,6 +7970,7 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_index( request, @@ -7866,6 +7982,7 @@ async def test_delete_index_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7963,10 +8080,14 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_upsert_datapoints" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, + "post_upsert_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_upsert_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.UpsertDatapointsRequest.pb( index_service.UpsertDatapointsRequest() ) @@ -7992,6 +8113,10 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.UpsertDatapointsResponse() + post_with_metadata.return_value = ( + index_service.UpsertDatapointsResponse(), + metadata, + ) await client.upsert_datapoints( request, @@ -8003,6 +8128,7 @@ async def test_upsert_datapoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8100,10 +8226,14 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "post_remove_datapoints" ) as post, mock.patch.object( + transports.AsyncIndexServiceRestInterceptor, + "post_remove_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncIndexServiceRestInterceptor, "pre_remove_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = index_service.RemoveDatapointsRequest.pb( index_service.RemoveDatapointsRequest() ) @@ -8129,6 +8259,10 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = index_service.RemoveDatapointsResponse() + post_with_metadata.return_value = ( + index_service.RemoveDatapointsResponse(), + metadata, + ) await client.remove_datapoints( request, @@ -8140,6 +8274,7 @@ async def test_remove_datapoints_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_job_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_job_service.py index 5e70239ad2..b3cdb18ca5 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_job_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_job_service.py @@ -127,6 +127,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -361,6 +369,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = JobServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = JobServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -24543,10 +24594,13 @@ def test_create_custom_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_create_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateCustomJobRequest.pb( job_service.CreateCustomJobRequest() ) @@ -24570,6 +24624,7 @@ def test_create_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_custom_job.CustomJob() + post_with_metadata.return_value = gca_custom_job.CustomJob(), metadata client.create_custom_job( request, @@ -24581,6 +24636,7 @@ def test_create_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_custom_job_rest_bad_request(request_type=job_service.GetCustomJobRequest): @@ -24671,10 +24727,13 @@ def test_get_custom_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetCustomJobRequest.pb( job_service.GetCustomJobRequest() ) @@ -24698,6 +24757,7 @@ def test_get_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = custom_job.CustomJob() + post_with_metadata.return_value = custom_job.CustomJob(), metadata client.get_custom_job( request, @@ -24709,6 +24769,7 @@ def test_get_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_custom_jobs_rest_bad_request( @@ -24793,10 +24854,13 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_custom_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_list_custom_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_custom_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListCustomJobsRequest.pb( job_service.ListCustomJobsRequest() ) @@ -24822,6 +24886,7 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListCustomJobsResponse() + post_with_metadata.return_value = job_service.ListCustomJobsResponse(), metadata client.list_custom_jobs( request, @@ -24833,6 +24898,7 @@ def test_list_custom_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_custom_job_rest_bad_request( @@ -24913,10 +24979,13 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_custom_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_delete_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteCustomJobRequest.pb( job_service.DeleteCustomJobRequest() ) @@ -24940,6 +25009,7 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_custom_job( request, @@ -24951,6 +25021,7 @@ def test_delete_custom_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_custom_job_rest_bad_request( @@ -25278,10 +25349,14 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateDataLabelingJobRequest.pb( job_service.CreateDataLabelingJobRequest() ) @@ -25307,6 +25382,10 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = ( + gca_data_labeling_job.DataLabelingJob(), + metadata, + ) client.create_data_labeling_job( request, @@ -25318,6 +25397,7 @@ def test_create_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_data_labeling_job_rest_bad_request( @@ -25422,10 +25502,13 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_data_labeling_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetDataLabelingJobRequest.pb( job_service.GetDataLabelingJobRequest() ) @@ -25451,6 +25534,7 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = data_labeling_job.DataLabelingJob(), metadata client.get_data_labeling_job( request, @@ -25462,6 +25546,7 @@ def test_get_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_data_labeling_jobs_rest_bad_request( @@ -25546,10 +25631,14 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_data_labeling_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_data_labeling_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_data_labeling_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListDataLabelingJobsRequest.pb( job_service.ListDataLabelingJobsRequest() ) @@ -25575,6 +25664,10 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListDataLabelingJobsResponse() + post_with_metadata.return_value = ( + job_service.ListDataLabelingJobsResponse(), + metadata, + ) client.list_data_labeling_jobs( request, @@ -25586,6 +25679,7 @@ def test_list_data_labeling_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_data_labeling_job_rest_bad_request( @@ -25670,10 +25764,14 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_data_labeling_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteDataLabelingJobRequest.pb( job_service.DeleteDataLabelingJobRequest() ) @@ -25697,6 +25795,7 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_data_labeling_job( request, @@ -25708,6 +25807,7 @@ def test_delete_data_labeling_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_data_labeling_job_rest_bad_request( @@ -26202,10 +26302,14 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateHyperparameterTuningJobRequest.pb( job_service.CreateHyperparameterTuningJobRequest() ) @@ -26231,6 +26335,10 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + gca_hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) client.create_hyperparameter_tuning_job( request, @@ -26242,6 +26350,7 @@ def test_create_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_hyperparameter_tuning_job_rest_bad_request( @@ -26346,10 +26455,14 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetHyperparameterTuningJobRequest.pb( job_service.GetHyperparameterTuningJobRequest() ) @@ -26375,6 +26488,10 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) client.get_hyperparameter_tuning_job( request, @@ -26386,6 +26503,7 @@ def test_get_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_hyperparameter_tuning_jobs_rest_bad_request( @@ -26470,10 +26588,14 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_hyperparameter_tuning_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_hyperparameter_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_hyperparameter_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListHyperparameterTuningJobsRequest.pb( job_service.ListHyperparameterTuningJobsRequest() ) @@ -26499,6 +26621,10 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListHyperparameterTuningJobsResponse() + post_with_metadata.return_value = ( + job_service.ListHyperparameterTuningJobsResponse(), + metadata, + ) client.list_hyperparameter_tuning_jobs( request, @@ -26510,6 +26636,7 @@ def test_list_hyperparameter_tuning_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_hyperparameter_tuning_job_rest_bad_request( @@ -26594,10 +26721,14 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteHyperparameterTuningJobRequest.pb( job_service.DeleteHyperparameterTuningJobRequest() ) @@ -26621,6 +26752,7 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_hyperparameter_tuning_job( request, @@ -26632,6 +26764,7 @@ def test_delete_hyperparameter_tuning_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_hyperparameter_tuning_job_rest_bad_request( @@ -27040,10 +27173,13 @@ def test_create_nas_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_create_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateNasJobRequest.pb( job_service.CreateNasJobRequest() ) @@ -27067,6 +27203,7 @@ def test_create_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_nas_job.NasJob() + post_with_metadata.return_value = gca_nas_job.NasJob(), metadata client.create_nas_job( request, @@ -27078,6 +27215,7 @@ def test_create_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_nas_job_rest_bad_request(request_type=job_service.GetNasJobRequest): @@ -27170,10 +27308,13 @@ def test_get_nas_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasJobRequest.pb(job_service.GetNasJobRequest()) transcode.return_value = { "method": "post", @@ -27195,6 +27336,7 @@ def test_get_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasJob() + post_with_metadata.return_value = nas_job.NasJob(), metadata client.get_nas_job( request, @@ -27206,6 +27348,7 @@ def test_get_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_nas_jobs_rest_bad_request(request_type=job_service.ListNasJobsRequest): @@ -27288,10 +27431,13 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_nas_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_list_nas_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_nas_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasJobsRequest.pb(job_service.ListNasJobsRequest()) transcode.return_value = { "method": "post", @@ -27315,6 +27461,7 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasJobsResponse() + post_with_metadata.return_value = job_service.ListNasJobsResponse(), metadata client.list_nas_jobs( request, @@ -27326,6 +27473,7 @@ def test_list_nas_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_nas_job_rest_bad_request(request_type=job_service.DeleteNasJobRequest): @@ -27404,10 +27552,13 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_nas_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_delete_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteNasJobRequest.pb( job_service.DeleteNasJobRequest() ) @@ -27431,6 +27582,7 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_nas_job( request, @@ -27442,6 +27594,7 @@ def test_delete_nas_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_nas_job_rest_bad_request(request_type=job_service.CancelNasJobRequest): @@ -27639,10 +27792,13 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_nas_trial_detail" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, "post_get_nas_trial_detail_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_nas_trial_detail" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasTrialDetailRequest.pb( job_service.GetNasTrialDetailRequest() ) @@ -27666,6 +27822,7 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasTrialDetail() + post_with_metadata.return_value = nas_job.NasTrialDetail(), metadata client.get_nas_trial_detail( request, @@ -27677,6 +27834,7 @@ def test_get_nas_trial_detail_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_nas_trial_details_rest_bad_request( @@ -27761,10 +27919,14 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_nas_trial_details" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_nas_trial_details_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_nas_trial_details" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasTrialDetailsRequest.pb( job_service.ListNasTrialDetailsRequest() ) @@ -27790,6 +27952,10 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasTrialDetailsResponse() + post_with_metadata.return_value = ( + job_service.ListNasTrialDetailsResponse(), + metadata, + ) client.list_nas_trial_details( request, @@ -27801,6 +27967,7 @@ def test_list_nas_trial_details_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_batch_prediction_job_rest_bad_request( @@ -27867,10 +28034,25 @@ def test_create_batch_prediction_job_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, }, "input_config": { @@ -28187,10 +28369,14 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_create_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateBatchPredictionJobRequest.pb( job_service.CreateBatchPredictionJobRequest() ) @@ -28216,6 +28402,10 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + gca_batch_prediction_job.BatchPredictionJob(), + metadata, + ) client.create_batch_prediction_job( request, @@ -28227,6 +28417,7 @@ def test_create_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_batch_prediction_job_rest_bad_request( @@ -28333,10 +28524,14 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetBatchPredictionJobRequest.pb( job_service.GetBatchPredictionJobRequest() ) @@ -28362,6 +28557,10 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + batch_prediction_job.BatchPredictionJob(), + metadata, + ) client.get_batch_prediction_job( request, @@ -28373,6 +28572,7 @@ def test_get_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_batch_prediction_jobs_rest_bad_request( @@ -28457,10 +28657,14 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_list_batch_prediction_jobs" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_batch_prediction_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_batch_prediction_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListBatchPredictionJobsRequest.pb( job_service.ListBatchPredictionJobsRequest() ) @@ -28486,6 +28690,10 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListBatchPredictionJobsResponse() + post_with_metadata.return_value = ( + job_service.ListBatchPredictionJobsResponse(), + metadata, + ) client.list_batch_prediction_jobs( request, @@ -28497,6 +28705,7 @@ def test_list_batch_prediction_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_batch_prediction_job_rest_bad_request( @@ -28581,10 +28790,14 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.JobServiceRestInterceptor, "post_delete_batch_prediction_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteBatchPredictionJobRequest.pb( job_service.DeleteBatchPredictionJobRequest() ) @@ -28608,6 +28821,7 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_batch_prediction_job( request, @@ -28619,6 +28833,7 @@ def test_delete_batch_prediction_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_batch_prediction_job_rest_bad_request( @@ -29018,11 +29233,15 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_create_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_create_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_create_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateModelDeploymentMonitoringJobRequest.pb( job_service.CreateModelDeploymentMonitoringJobRequest() ) @@ -29052,6 +29271,10 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept post.return_value = ( gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) client.create_model_deployment_monitoring_job( request, @@ -29063,6 +29286,7 @@ def test_create_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_model_deployment_monitoring_stats_anomalies_rest_bad_request( @@ -29164,11 +29388,15 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( transports.JobServiceRestInterceptor, "post_search_model_deployment_monitoring_stats_anomalies", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_search_model_deployment_monitoring_stats_anomalies_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_search_model_deployment_monitoring_stats_anomalies", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest.pb( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest() @@ -29200,6 +29428,10 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( post.return_value = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse() ) + post_with_metadata.return_value = ( + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse(), + metadata, + ) client.search_model_deployment_monitoring_stats_anomalies( request, @@ -29211,6 +29443,7 @@ def test_search_model_deployment_monitoring_stats_anomalies_rest_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_deployment_monitoring_job_rest_bad_request( @@ -29324,10 +29557,14 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.JobServiceRestInterceptor, "post_get_model_deployment_monitoring_job" ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_get_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_get_model_deployment_monitoring_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetModelDeploymentMonitoringJobRequest.pb( job_service.GetModelDeploymentMonitoringJobRequest() ) @@ -29357,6 +29594,10 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) post.return_value = ( model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) client.get_model_deployment_monitoring_job( request, @@ -29368,6 +29609,7 @@ def test_get_model_deployment_monitoring_job_rest_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_deployment_monitoring_jobs_rest_bad_request( @@ -29455,11 +29697,15 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto transports.JobServiceRestInterceptor, "post_list_model_deployment_monitoring_jobs", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_list_model_deployment_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_list_model_deployment_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListModelDeploymentMonitoringJobsRequest.pb( job_service.ListModelDeploymentMonitoringJobsRequest() ) @@ -29485,6 +29731,10 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = job_service.ListModelDeploymentMonitoringJobsResponse() + post_with_metadata.return_value = ( + job_service.ListModelDeploymentMonitoringJobsResponse(), + metadata, + ) client.list_model_deployment_monitoring_jobs( request, @@ -29496,6 +29746,7 @@ def test_list_model_deployment_monitoring_jobs_rest_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_model_deployment_monitoring_job_rest_bad_request( @@ -29759,11 +30010,15 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_update_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_update_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_update_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.UpdateModelDeploymentMonitoringJobRequest.pb( job_service.UpdateModelDeploymentMonitoringJobRequest() ) @@ -29787,6 +30042,7 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_model_deployment_monitoring_job( request, @@ -29798,6 +30054,7 @@ def test_update_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_deployment_monitoring_job_rest_bad_request( @@ -29883,11 +30140,15 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept transports.JobServiceRestInterceptor, "post_delete_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.JobServiceRestInterceptor, + "post_delete_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.JobServiceRestInterceptor, "pre_delete_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteModelDeploymentMonitoringJobRequest.pb( job_service.DeleteModelDeploymentMonitoringJobRequest() ) @@ -29911,6 +30172,7 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_deployment_monitoring_job( request, @@ -29922,6 +30184,7 @@ def test_delete_model_deployment_monitoring_job_rest_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_pause_model_deployment_monitoring_job_rest_bad_request( @@ -31831,10 +32094,14 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_custom_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateCustomJobRequest.pb( job_service.CreateCustomJobRequest() ) @@ -31858,6 +32125,7 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_custom_job.CustomJob() + post_with_metadata.return_value = gca_custom_job.CustomJob(), metadata await client.create_custom_job( request, @@ -31869,6 +32137,7 @@ async def test_create_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31977,10 +32246,13 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_get_custom_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetCustomJobRequest.pb( job_service.GetCustomJobRequest() ) @@ -32004,6 +32276,7 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = custom_job.CustomJob() + post_with_metadata.return_value = custom_job.CustomJob(), metadata await client.get_custom_job( request, @@ -32015,6 +32288,7 @@ async def test_get_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32115,10 +32389,13 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_custom_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_list_custom_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_custom_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListCustomJobsRequest.pb( job_service.ListCustomJobsRequest() ) @@ -32144,6 +32421,7 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListCustomJobsResponse() + post_with_metadata.return_value = job_service.ListCustomJobsResponse(), metadata await client.list_custom_jobs( request, @@ -32155,6 +32433,7 @@ async def test_list_custom_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32251,10 +32530,14 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_custom_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_custom_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_custom_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteCustomJobRequest.pb( job_service.DeleteCustomJobRequest() ) @@ -32278,6 +32561,7 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_custom_job( request, @@ -32289,6 +32573,7 @@ async def test_delete_custom_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32648,10 +32933,14 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateDataLabelingJobRequest.pb( job_service.CreateDataLabelingJobRequest() ) @@ -32677,6 +32966,10 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = gca_data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = ( + gca_data_labeling_job.DataLabelingJob(), + metadata, + ) await client.create_data_labeling_job( request, @@ -32688,6 +32981,7 @@ async def test_create_data_labeling_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32808,10 +33102,14 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetDataLabelingJobRequest.pb( job_service.GetDataLabelingJobRequest() ) @@ -32837,6 +33135,7 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = data_labeling_job.DataLabelingJob() + post_with_metadata.return_value = data_labeling_job.DataLabelingJob(), metadata await client.get_data_labeling_job( request, @@ -32848,6 +33147,7 @@ async def test_get_data_labeling_job_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32948,10 +33248,14 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_data_labeling_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_data_labeling_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_data_labeling_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListDataLabelingJobsRequest.pb( job_service.ListDataLabelingJobsRequest() ) @@ -32977,6 +33281,10 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = job_service.ListDataLabelingJobsResponse() + post_with_metadata.return_value = ( + job_service.ListDataLabelingJobsResponse(), + metadata, + ) await client.list_data_labeling_jobs( request, @@ -32988,6 +33296,7 @@ async def test_list_data_labeling_jobs_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33088,10 +33397,14 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_data_labeling_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_data_labeling_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_data_labeling_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteDataLabelingJobRequest.pb( job_service.DeleteDataLabelingJobRequest() ) @@ -33115,6 +33428,7 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_data_labeling_job( request, @@ -33126,6 +33440,7 @@ async def test_delete_data_labeling_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33655,11 +33970,15 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_create_hyperparameter_tuning_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_hyperparameter_tuning_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateHyperparameterTuningJobRequest.pb( job_service.CreateHyperparameterTuningJobRequest() ) @@ -33685,6 +34004,10 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + gca_hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) await client.create_hyperparameter_tuning_job( request, @@ -33696,6 +34019,7 @@ async def test_create_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33818,10 +34142,14 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_hyperparameter_tuning_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_hyperparameter_tuning_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetHyperparameterTuningJobRequest.pb( job_service.GetHyperparameterTuningJobRequest() ) @@ -33847,6 +34175,10 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = hyperparameter_tuning_job.HyperparameterTuningJob() + post_with_metadata.return_value = ( + hyperparameter_tuning_job.HyperparameterTuningJob(), + metadata, + ) await client.get_hyperparameter_tuning_job( request, @@ -33858,6 +34190,7 @@ async def test_get_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -33961,10 +34294,14 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_list_hyperparameter_tuning_jobs", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_hyperparameter_tuning_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_hyperparameter_tuning_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListHyperparameterTuningJobsRequest.pb( job_service.ListHyperparameterTuningJobsRequest() ) @@ -33990,6 +34327,10 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = job_service.ListHyperparameterTuningJobsResponse() + post_with_metadata.return_value = ( + job_service.ListHyperparameterTuningJobsResponse(), + metadata, + ) await client.list_hyperparameter_tuning_jobs( request, @@ -34001,6 +34342,7 @@ async def test_list_hyperparameter_tuning_jobs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34104,11 +34446,15 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_delete_hyperparameter_tuning_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_hyperparameter_tuning_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_hyperparameter_tuning_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteHyperparameterTuningJobRequest.pb( job_service.DeleteHyperparameterTuningJobRequest() ) @@ -34132,6 +34478,7 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_hyperparameter_tuning_job( request, @@ -34143,6 +34490,7 @@ async def test_delete_hyperparameter_tuning_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34588,10 +34936,13 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_create_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateNasJobRequest.pb( job_service.CreateNasJobRequest() ) @@ -34615,6 +34966,7 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_nas_job.NasJob() + post_with_metadata.return_value = gca_nas_job.NasJob(), metadata await client.create_nas_job( request, @@ -34626,6 +34978,7 @@ async def test_create_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34736,10 +35089,13 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_get_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasJobRequest.pb(job_service.GetNasJobRequest()) transcode.return_value = { "method": "post", @@ -34761,6 +35117,7 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasJob() + post_with_metadata.return_value = nas_job.NasJob(), metadata await client.get_nas_job( request, @@ -34772,6 +35129,7 @@ async def test_get_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -34872,10 +35230,13 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_nas_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_list_nas_jobs_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_nas_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasJobsRequest.pb(job_service.ListNasJobsRequest()) transcode.return_value = { "method": "post", @@ -34899,6 +35260,7 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = job_service.ListNasJobsResponse() + post_with_metadata.return_value = job_service.ListNasJobsResponse(), metadata await client.list_nas_jobs( request, @@ -34910,6 +35272,7 @@ async def test_list_nas_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35006,10 +35369,13 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_nas_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, "post_delete_nas_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_nas_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteNasJobRequest.pb( job_service.DeleteNasJobRequest() ) @@ -35033,6 +35399,7 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_nas_job( request, @@ -35044,6 +35411,7 @@ async def test_delete_nas_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35275,10 +35643,14 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_nas_trial_detail" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_nas_trial_detail_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_nas_trial_detail" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetNasTrialDetailRequest.pb( job_service.GetNasTrialDetailRequest() ) @@ -35302,6 +35674,7 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = nas_job.NasTrialDetail() + post_with_metadata.return_value = nas_job.NasTrialDetail(), metadata await client.get_nas_trial_detail( request, @@ -35313,6 +35686,7 @@ async def test_get_nas_trial_detail_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35413,10 +35787,14 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_nas_trial_details" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_nas_trial_details_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_nas_trial_details" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListNasTrialDetailsRequest.pb( job_service.ListNasTrialDetailsRequest() ) @@ -35442,6 +35820,10 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = job_service.ListNasTrialDetailsResponse() + post_with_metadata.return_value = ( + job_service.ListNasTrialDetailsResponse(), + metadata, + ) await client.list_nas_trial_details( request, @@ -35453,6 +35835,7 @@ async def test_list_nas_trial_details_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -35528,10 +35911,25 @@ async def test_create_batch_prediction_job_rest_asyncio_call_success(request_typ "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, }, "input_config": { @@ -35855,10 +36253,14 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_create_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateBatchPredictionJobRequest.pb( job_service.CreateBatchPredictionJobRequest() ) @@ -35884,6 +36286,10 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = gca_batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + gca_batch_prediction_job.BatchPredictionJob(), + metadata, + ) await client.create_batch_prediction_job( request, @@ -35895,6 +36301,7 @@ async def test_create_batch_prediction_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36017,10 +36424,14 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_get_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetBatchPredictionJobRequest.pb( job_service.GetBatchPredictionJobRequest() ) @@ -36046,6 +36457,10 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = batch_prediction_job.BatchPredictionJob() + post_with_metadata.return_value = ( + batch_prediction_job.BatchPredictionJob(), + metadata, + ) await client.get_batch_prediction_job( request, @@ -36057,6 +36472,7 @@ async def test_get_batch_prediction_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36157,10 +36573,14 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce ) as transcode, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_list_batch_prediction_jobs" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_batch_prediction_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_batch_prediction_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListBatchPredictionJobsRequest.pb( job_service.ListBatchPredictionJobsRequest() ) @@ -36186,6 +36606,10 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = job_service.ListBatchPredictionJobsResponse() + post_with_metadata.return_value = ( + job_service.ListBatchPredictionJobsResponse(), + metadata, + ) await client.list_batch_prediction_jobs( request, @@ -36197,6 +36621,7 @@ async def test_list_batch_prediction_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36297,10 +36722,14 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc ), mock.patch.object( transports.AsyncJobServiceRestInterceptor, "post_delete_batch_prediction_job" ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_batch_prediction_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_batch_prediction_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteBatchPredictionJobRequest.pb( job_service.DeleteBatchPredictionJobRequest() ) @@ -36324,6 +36753,7 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_batch_prediction_job( request, @@ -36335,6 +36765,7 @@ async def test_delete_batch_prediction_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36770,11 +37201,15 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_create_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_create_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_create_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.CreateModelDeploymentMonitoringJobRequest.pb( job_service.CreateModelDeploymentMonitoringJobRequest() ) @@ -36804,6 +37239,10 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( post.return_value = ( gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + gca_model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) await client.create_model_deployment_monitoring_job( request, @@ -36815,6 +37254,7 @@ async def test_create_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -36934,11 +37374,15 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i transports.AsyncJobServiceRestInterceptor, "post_search_model_deployment_monitoring_stats_anomalies", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_search_model_deployment_monitoring_stats_anomalies_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_search_model_deployment_monitoring_stats_anomalies", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest.pb( job_service.SearchModelDeploymentMonitoringStatsAnomaliesRequest() @@ -36970,6 +37414,10 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i post.return_value = ( job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse() ) + post_with_metadata.return_value = ( + job_service.SearchModelDeploymentMonitoringStatsAnomaliesResponse(), + metadata, + ) await client.search_model_deployment_monitoring_stats_anomalies( request, @@ -36981,6 +37429,7 @@ async def test_search_model_deployment_monitoring_stats_anomalies_rest_asyncio_i pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37115,11 +37564,15 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_get_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_get_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_get_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.GetModelDeploymentMonitoringJobRequest.pb( job_service.GetModelDeploymentMonitoringJobRequest() ) @@ -37149,6 +37602,10 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( post.return_value = ( model_deployment_monitoring_job.ModelDeploymentMonitoringJob() ) + post_with_metadata.return_value = ( + model_deployment_monitoring_job.ModelDeploymentMonitoringJob(), + metadata, + ) await client.get_model_deployment_monitoring_job( request, @@ -37160,6 +37617,7 @@ async def test_get_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37267,11 +37725,15 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_list_model_deployment_monitoring_jobs", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_list_model_deployment_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_list_model_deployment_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.ListModelDeploymentMonitoringJobsRequest.pb( job_service.ListModelDeploymentMonitoringJobsRequest() ) @@ -37297,6 +37759,10 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = job_service.ListModelDeploymentMonitoringJobsResponse() + post_with_metadata.return_value = ( + job_service.ListModelDeploymentMonitoringJobsResponse(), + metadata, + ) await client.list_model_deployment_monitoring_jobs( request, @@ -37308,6 +37774,7 @@ async def test_list_model_deployment_monitoring_jobs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37591,11 +38058,15 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_update_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_update_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_update_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.UpdateModelDeploymentMonitoringJobRequest.pb( job_service.UpdateModelDeploymentMonitoringJobRequest() ) @@ -37619,6 +38090,7 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_model_deployment_monitoring_job( request, @@ -37630,6 +38102,7 @@ async def test_update_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -37735,11 +38208,15 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( transports.AsyncJobServiceRestInterceptor, "post_delete_model_deployment_monitoring_job", ) as post, mock.patch.object( + transports.AsyncJobServiceRestInterceptor, + "post_delete_model_deployment_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncJobServiceRestInterceptor, "pre_delete_model_deployment_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = job_service.DeleteModelDeploymentMonitoringJobRequest.pb( job_service.DeleteModelDeploymentMonitoringJobRequest() ) @@ -37763,6 +38240,7 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_deployment_monitoring_job( request, @@ -37774,6 +38252,7 @@ async def test_delete_model_deployment_monitoring_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_llm_utility_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_llm_utility_service.py index dbda418ab7..12f90c3a38 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_llm_utility_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_llm_utility_service.py @@ -83,6 +83,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -349,6 +357,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = LlmUtilityServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = LlmUtilityServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1885,10 +1936,13 @@ def test_compute_tokens_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "post_compute_tokens" ) as post, mock.patch.object( + transports.LlmUtilityServiceRestInterceptor, "post_compute_tokens_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.LlmUtilityServiceRestInterceptor, "pre_compute_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = llm_utility_service.ComputeTokensRequest.pb( llm_utility_service.ComputeTokensRequest() ) @@ -1914,6 +1968,10 @@ def test_compute_tokens_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = llm_utility_service.ComputeTokensResponse() + post_with_metadata.return_value = ( + llm_utility_service.ComputeTokensResponse(), + metadata, + ) client.compute_tokens( request, @@ -1925,6 +1983,7 @@ def test_compute_tokens_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -2685,10 +2744,14 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "post_compute_tokens" ) as post, mock.patch.object( + transports.AsyncLlmUtilityServiceRestInterceptor, + "post_compute_tokens_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncLlmUtilityServiceRestInterceptor, "pre_compute_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = llm_utility_service.ComputeTokensRequest.pb( llm_utility_service.ComputeTokensRequest() ) @@ -2714,6 +2777,10 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = llm_utility_service.ComputeTokensResponse() + post_with_metadata.return_value = ( + llm_utility_service.ComputeTokensResponse(), + metadata, + ) await client.compute_tokens( request, @@ -2725,6 +2792,7 @@ async def test_compute_tokens_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_match_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_match_service.py index 85353fcbb8..dd34b58d9a 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_match_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_match_service.py @@ -78,6 +78,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -320,6 +328,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MatchServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MatchServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2097,10 +2148,13 @@ def test_find_neighbors_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MatchServiceRestInterceptor, "post_find_neighbors" ) as post, mock.patch.object( + transports.MatchServiceRestInterceptor, "post_find_neighbors_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MatchServiceRestInterceptor, "pre_find_neighbors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.FindNeighborsRequest.pb( match_service.FindNeighborsRequest() ) @@ -2126,6 +2180,10 @@ def test_find_neighbors_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.FindNeighborsResponse() + post_with_metadata.return_value = ( + match_service.FindNeighborsResponse(), + metadata, + ) client.find_neighbors( request, @@ -2137,6 +2195,7 @@ def test_find_neighbors_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_index_datapoints_rest_bad_request( @@ -2222,10 +2281,14 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MatchServiceRestInterceptor, "post_read_index_datapoints" ) as post, mock.patch.object( + transports.MatchServiceRestInterceptor, + "post_read_index_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MatchServiceRestInterceptor, "pre_read_index_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.ReadIndexDatapointsRequest.pb( match_service.ReadIndexDatapointsRequest() ) @@ -2251,6 +2314,10 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.ReadIndexDatapointsResponse() + post_with_metadata.return_value = ( + match_service.ReadIndexDatapointsResponse(), + metadata, + ) client.read_index_datapoints( request, @@ -2262,6 +2329,7 @@ def test_read_index_datapoints_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3048,10 +3116,13 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "post_find_neighbors" ) as post, mock.patch.object( + transports.AsyncMatchServiceRestInterceptor, "post_find_neighbors_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "pre_find_neighbors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.FindNeighborsRequest.pb( match_service.FindNeighborsRequest() ) @@ -3077,6 +3148,10 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = match_service.FindNeighborsResponse() + post_with_metadata.return_value = ( + match_service.FindNeighborsResponse(), + metadata, + ) await client.find_neighbors( request, @@ -3088,6 +3163,7 @@ async def test_find_neighbors_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3189,10 +3265,14 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "post_read_index_datapoints" ) as post, mock.patch.object( + transports.AsyncMatchServiceRestInterceptor, + "post_read_index_datapoints_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMatchServiceRestInterceptor, "pre_read_index_datapoints" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = match_service.ReadIndexDatapointsRequest.pb( match_service.ReadIndexDatapointsRequest() ) @@ -3218,6 +3298,10 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = match_service.ReadIndexDatapointsResponse() + post_with_metadata.return_value = ( + match_service.ReadIndexDatapointsResponse(), + metadata, + ) await client.read_index_datapoints( request, @@ -3229,6 +3313,7 @@ async def test_read_index_datapoints_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_metadata_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_metadata_service.py index 744ef459a2..9b2ec92b42 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_metadata_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_metadata_service.py @@ -102,6 +102,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -359,6 +367,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MetadataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MetadataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -21571,10 +21622,14 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_create_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataStoreRequest.pb( metadata_service.CreateMetadataStoreRequest() ) @@ -21598,6 +21653,7 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_metadata_store( request, @@ -21609,6 +21665,7 @@ def test_create_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_metadata_store_rest_bad_request( @@ -21695,10 +21752,14 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_get_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataStoreRequest.pb( metadata_service.GetMetadataStoreRequest() ) @@ -21724,6 +21785,7 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_store.MetadataStore() + post_with_metadata.return_value = metadata_store.MetadataStore(), metadata client.get_metadata_store( request, @@ -21735,6 +21797,7 @@ def test_get_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_metadata_stores_rest_bad_request( @@ -21819,10 +21882,14 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_metadata_stores" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_list_metadata_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_metadata_stores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataStoresRequest.pb( metadata_service.ListMetadataStoresRequest() ) @@ -21848,6 +21915,10 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataStoresResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataStoresResponse(), + metadata, + ) client.list_metadata_stores( request, @@ -21859,6 +21930,7 @@ def test_list_metadata_stores_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_metadata_store_rest_bad_request( @@ -21939,10 +22011,14 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_metadata_store" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_delete_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteMetadataStoreRequest.pb( metadata_service.DeleteMetadataStoreRequest() ) @@ -21966,6 +22042,7 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_metadata_store( request, @@ -21977,6 +22054,7 @@ def test_delete_metadata_store_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_artifact_rest_bad_request( @@ -22160,10 +22238,13 @@ def test_create_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateArtifactRequest.pb( metadata_service.CreateArtifactRequest() ) @@ -22187,6 +22268,7 @@ def test_create_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata client.create_artifact( request, @@ -22198,6 +22280,7 @@ def test_create_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_artifact_rest_bad_request( @@ -22300,10 +22383,13 @@ def test_get_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetArtifactRequest.pb( metadata_service.GetArtifactRequest() ) @@ -22327,6 +22413,7 @@ def test_get_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = artifact.Artifact() + post_with_metadata.return_value = artifact.Artifact(), metadata client.get_artifact( request, @@ -22338,6 +22425,7 @@ def test_get_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_artifacts_rest_bad_request( @@ -22426,10 +22514,13 @@ def test_list_artifacts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_artifacts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_artifacts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListArtifactsRequest.pb( metadata_service.ListArtifactsRequest() ) @@ -22455,6 +22546,10 @@ def test_list_artifacts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListArtifactsResponse() + post_with_metadata.return_value = ( + metadata_service.ListArtifactsResponse(), + metadata, + ) client.list_artifacts( request, @@ -22466,6 +22561,7 @@ def test_list_artifacts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_artifact_rest_bad_request( @@ -22653,10 +22749,13 @@ def test_update_artifact_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateArtifactRequest.pb( metadata_service.UpdateArtifactRequest() ) @@ -22680,6 +22779,7 @@ def test_update_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata client.update_artifact( request, @@ -22691,6 +22791,7 @@ def test_update_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_artifact_rest_bad_request( @@ -22775,10 +22876,13 @@ def test_delete_artifact_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_artifact" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_artifact_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteArtifactRequest.pb( metadata_service.DeleteArtifactRequest() ) @@ -22802,6 +22906,7 @@ def test_delete_artifact_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_artifact( request, @@ -22813,6 +22918,7 @@ def test_delete_artifact_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_artifacts_rest_bad_request( @@ -22897,10 +23003,13 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_artifacts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_artifacts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeArtifactsRequest.pb( metadata_service.PurgeArtifactsRequest() ) @@ -22924,6 +23033,7 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_artifacts( request, @@ -22935,6 +23045,7 @@ def test_purge_artifacts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_context_rest_bad_request( @@ -23115,10 +23226,13 @@ def test_create_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateContextRequest.pb( metadata_service.CreateContextRequest() ) @@ -23142,6 +23256,7 @@ def test_create_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata client.create_context( request, @@ -23153,6 +23268,7 @@ def test_create_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_context_rest_bad_request(request_type=metadata_service.GetContextRequest): @@ -23251,10 +23367,13 @@ def test_get_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetContextRequest.pb( metadata_service.GetContextRequest() ) @@ -23278,6 +23397,7 @@ def test_get_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = context.Context() + post_with_metadata.return_value = context.Context(), metadata client.get_context( request, @@ -23289,6 +23409,7 @@ def test_get_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_contexts_rest_bad_request( @@ -23377,10 +23498,13 @@ def test_list_contexts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_contexts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_contexts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListContextsRequest.pb( metadata_service.ListContextsRequest() ) @@ -23406,6 +23530,10 @@ def test_list_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListContextsResponse() + post_with_metadata.return_value = ( + metadata_service.ListContextsResponse(), + metadata, + ) client.list_contexts( request, @@ -23417,6 +23545,7 @@ def test_list_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_context_rest_bad_request( @@ -23601,10 +23730,13 @@ def test_update_context_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateContextRequest.pb( metadata_service.UpdateContextRequest() ) @@ -23628,6 +23760,7 @@ def test_update_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata client.update_context( request, @@ -23639,6 +23772,7 @@ def test_update_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_context_rest_bad_request( @@ -23723,10 +23857,13 @@ def test_delete_context_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_context" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteContextRequest.pb( metadata_service.DeleteContextRequest() ) @@ -23750,6 +23887,7 @@ def test_delete_context_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_context( request, @@ -23761,6 +23899,7 @@ def test_delete_context_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_contexts_rest_bad_request( @@ -23845,10 +23984,13 @@ def test_purge_contexts_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_contexts" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_contexts_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeContextsRequest.pb( metadata_service.PurgeContextsRequest() ) @@ -23872,6 +24014,7 @@ def test_purge_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_contexts( request, @@ -23883,6 +24026,7 @@ def test_purge_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_context_artifacts_and_executions_rest_bad_request( @@ -23973,11 +24117,15 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor transports.MetadataServiceRestInterceptor, "post_add_context_artifacts_and_executions", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_context_artifacts_and_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_context_artifacts_and_executions", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextArtifactsAndExecutionsRequest.pb( metadata_service.AddContextArtifactsAndExecutionsRequest() ) @@ -24005,6 +24153,10 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextArtifactsAndExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextArtifactsAndExecutionsResponse(), + metadata, + ) client.add_context_artifacts_and_executions( request, @@ -24016,6 +24168,7 @@ def test_add_context_artifacts_and_executions_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_context_children_rest_bad_request( @@ -24101,10 +24254,14 @@ def test_add_context_children_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_add_context_children" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextChildrenRequest.pb( metadata_service.AddContextChildrenRequest() ) @@ -24130,6 +24287,10 @@ def test_add_context_children_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextChildrenResponse(), + metadata, + ) client.add_context_children( request, @@ -24141,6 +24302,7 @@ def test_add_context_children_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_remove_context_children_rest_bad_request( @@ -24226,10 +24388,14 @@ def test_remove_context_children_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_remove_context_children" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_remove_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_remove_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.RemoveContextChildrenRequest.pb( metadata_service.RemoveContextChildrenRequest() ) @@ -24255,6 +24421,10 @@ def test_remove_context_children_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.RemoveContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.RemoveContextChildrenResponse(), + metadata, + ) client.remove_context_children( request, @@ -24266,6 +24436,7 @@ def test_remove_context_children_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_context_lineage_subgraph_rest_bad_request( @@ -24351,10 +24522,14 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_query_context_lineage_subgraph" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_context_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_context_lineage_subgraph" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryContextLineageSubgraphRequest.pb( metadata_service.QueryContextLineageSubgraphRequest() ) @@ -24380,6 +24555,7 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_context_lineage_subgraph( request, @@ -24391,6 +24567,7 @@ def test_query_context_lineage_subgraph_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_execution_rest_bad_request( @@ -24571,10 +24748,13 @@ def test_create_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_create_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateExecutionRequest.pb( metadata_service.CreateExecutionRequest() ) @@ -24598,6 +24778,7 @@ def test_create_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata client.create_execution( request, @@ -24609,6 +24790,7 @@ def test_create_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_execution_rest_bad_request( @@ -24709,10 +24891,13 @@ def test_get_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_get_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetExecutionRequest.pb( metadata_service.GetExecutionRequest() ) @@ -24736,6 +24921,7 @@ def test_get_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = execution.Execution() + post_with_metadata.return_value = execution.Execution(), metadata client.get_execution( request, @@ -24747,6 +24933,7 @@ def test_get_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_executions_rest_bad_request( @@ -24835,10 +25022,13 @@ def test_list_executions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_executions" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_list_executions_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListExecutionsRequest.pb( metadata_service.ListExecutionsRequest() ) @@ -24864,6 +25054,10 @@ def test_list_executions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.ListExecutionsResponse(), + metadata, + ) client.list_executions( request, @@ -24875,6 +25069,7 @@ def test_list_executions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_execution_rest_bad_request( @@ -25059,10 +25254,13 @@ def test_update_execution_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_update_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_update_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_update_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateExecutionRequest.pb( metadata_service.UpdateExecutionRequest() ) @@ -25086,6 +25284,7 @@ def test_update_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata client.update_execution( request, @@ -25097,6 +25296,7 @@ def test_update_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_execution_rest_bad_request( @@ -25181,10 +25381,13 @@ def test_delete_execution_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_delete_execution" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_delete_execution_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_delete_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteExecutionRequest.pb( metadata_service.DeleteExecutionRequest() ) @@ -25208,6 +25411,7 @@ def test_delete_execution_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_execution( request, @@ -25219,6 +25423,7 @@ def test_delete_execution_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_purge_executions_rest_bad_request( @@ -25303,10 +25508,13 @@ def test_purge_executions_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MetadataServiceRestInterceptor, "post_purge_executions" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, "post_purge_executions_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_purge_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeExecutionsRequest.pb( metadata_service.PurgeExecutionsRequest() ) @@ -25330,6 +25538,7 @@ def test_purge_executions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.purge_executions( request, @@ -25341,6 +25550,7 @@ def test_purge_executions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_execution_events_rest_bad_request( @@ -25426,10 +25636,14 @@ def test_add_execution_events_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_add_execution_events" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_add_execution_events_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_add_execution_events" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddExecutionEventsRequest.pb( metadata_service.AddExecutionEventsRequest() ) @@ -25455,6 +25669,10 @@ def test_add_execution_events_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddExecutionEventsResponse() + post_with_metadata.return_value = ( + metadata_service.AddExecutionEventsResponse(), + metadata, + ) client.add_execution_events( request, @@ -25466,6 +25684,7 @@ def test_add_execution_events_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_execution_inputs_and_outputs_rest_bad_request( @@ -25552,11 +25771,15 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): transports.MetadataServiceRestInterceptor, "post_query_execution_inputs_and_outputs", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_execution_inputs_and_outputs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_execution_inputs_and_outputs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryExecutionInputsAndOutputsRequest.pb( metadata_service.QueryExecutionInputsAndOutputsRequest() ) @@ -25582,6 +25805,7 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_execution_inputs_and_outputs( request, @@ -25593,6 +25817,7 @@ def test_query_execution_inputs_and_outputs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_metadata_schema_rest_bad_request( @@ -25769,10 +25994,14 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_create_metadata_schema" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_create_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_create_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataSchemaRequest.pb( metadata_service.CreateMetadataSchemaRequest() ) @@ -25798,6 +26027,7 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_metadata_schema.MetadataSchema() + post_with_metadata.return_value = gca_metadata_schema.MetadataSchema(), metadata client.create_metadata_schema( request, @@ -25809,6 +26039,7 @@ def test_create_metadata_schema_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_metadata_schema_rest_bad_request( @@ -25908,10 +26139,14 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_get_metadata_schema" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_get_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_get_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataSchemaRequest.pb( metadata_service.GetMetadataSchemaRequest() ) @@ -25937,6 +26172,7 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_schema.MetadataSchema() + post_with_metadata.return_value = metadata_schema.MetadataSchema(), metadata client.get_metadata_schema( request, @@ -25948,6 +26184,7 @@ def test_get_metadata_schema_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_metadata_schemas_rest_bad_request( @@ -26036,10 +26273,14 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MetadataServiceRestInterceptor, "post_list_metadata_schemas" ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_list_metadata_schemas_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_list_metadata_schemas" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataSchemasRequest.pb( metadata_service.ListMetadataSchemasRequest() ) @@ -26065,6 +26306,10 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataSchemasResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataSchemasResponse(), + metadata, + ) client.list_metadata_schemas( request, @@ -26076,6 +26321,7 @@ def test_list_metadata_schemas_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_query_artifact_lineage_subgraph_rest_bad_request( @@ -26162,10 +26408,14 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): transports.MetadataServiceRestInterceptor, "post_query_artifact_lineage_subgraph", ) as post, mock.patch.object( + transports.MetadataServiceRestInterceptor, + "post_query_artifact_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MetadataServiceRestInterceptor, "pre_query_artifact_lineage_subgraph" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryArtifactLineageSubgraphRequest.pb( metadata_service.QueryArtifactLineageSubgraphRequest() ) @@ -26191,6 +26441,7 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata client.query_artifact_lineage_subgraph( request, @@ -26202,6 +26453,7 @@ def test_query_artifact_lineage_subgraph_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -27704,10 +27956,14 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataStoreRequest.pb( metadata_service.CreateMetadataStoreRequest() ) @@ -27731,6 +27987,7 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_metadata_store( request, @@ -27742,6 +27999,7 @@ async def test_create_metadata_store_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27844,10 +28102,14 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataStoreRequest.pb( metadata_service.GetMetadataStoreRequest() ) @@ -27873,6 +28135,7 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_store.MetadataStore() + post_with_metadata.return_value = metadata_store.MetadataStore(), metadata await client.get_metadata_store( request, @@ -27884,6 +28147,7 @@ async def test_get_metadata_store_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27984,10 +28248,14 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_metadata_stores" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_metadata_stores_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_metadata_stores" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataStoresRequest.pb( metadata_service.ListMetadataStoresRequest() ) @@ -28013,6 +28281,10 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataStoresResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataStoresResponse(), + metadata, + ) await client.list_metadata_stores( request, @@ -28024,6 +28296,7 @@ async def test_list_metadata_stores_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28120,10 +28393,14 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_metadata_store" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_metadata_store_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_metadata_store" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteMetadataStoreRequest.pb( metadata_service.DeleteMetadataStoreRequest() ) @@ -28147,6 +28424,7 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_metadata_store( request, @@ -28158,6 +28436,7 @@ async def test_delete_metadata_store_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28357,10 +28636,14 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateArtifactRequest.pb( metadata_service.CreateArtifactRequest() ) @@ -28384,6 +28667,7 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata await client.create_artifact( request, @@ -28395,6 +28679,7 @@ async def test_create_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28513,10 +28798,14 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetArtifactRequest.pb( metadata_service.GetArtifactRequest() ) @@ -28540,6 +28829,7 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = artifact.Artifact() + post_with_metadata.return_value = artifact.Artifact(), metadata await client.get_artifact( request, @@ -28551,6 +28841,7 @@ async def test_get_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28655,10 +28946,14 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_artifacts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_artifacts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListArtifactsRequest.pb( metadata_service.ListArtifactsRequest() ) @@ -28684,6 +28979,10 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListArtifactsResponse() + post_with_metadata.return_value = ( + metadata_service.ListArtifactsResponse(), + metadata, + ) await client.list_artifacts( request, @@ -28695,6 +28994,7 @@ async def test_list_artifacts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28898,10 +29198,14 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateArtifactRequest.pb( metadata_service.UpdateArtifactRequest() ) @@ -28925,6 +29229,7 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_artifact.Artifact() + post_with_metadata.return_value = gca_artifact.Artifact(), metadata await client.update_artifact( request, @@ -28936,6 +29241,7 @@ async def test_update_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29036,10 +29342,14 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_artifact" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_artifact_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_artifact" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteArtifactRequest.pb( metadata_service.DeleteArtifactRequest() ) @@ -29063,6 +29373,7 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_artifact( request, @@ -29074,6 +29385,7 @@ async def test_delete_artifact_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29174,10 +29486,14 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_artifacts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_artifacts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_artifacts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeArtifactsRequest.pb( metadata_service.PurgeArtifactsRequest() ) @@ -29201,6 +29517,7 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_artifacts( request, @@ -29212,6 +29529,7 @@ async def test_purge_artifacts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29408,10 +29726,14 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateContextRequest.pb( metadata_service.CreateContextRequest() ) @@ -29435,6 +29757,7 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata await client.create_context( request, @@ -29446,6 +29769,7 @@ async def test_create_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29562,10 +29886,13 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, "post_get_context_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetContextRequest.pb( metadata_service.GetContextRequest() ) @@ -29589,6 +29916,7 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = context.Context() + post_with_metadata.return_value = context.Context(), metadata await client.get_context( request, @@ -29600,6 +29928,7 @@ async def test_get_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29704,10 +30033,14 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_contexts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListContextsRequest.pb( metadata_service.ListContextsRequest() ) @@ -29733,6 +30066,10 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListContextsResponse() + post_with_metadata.return_value = ( + metadata_service.ListContextsResponse(), + metadata, + ) await client.list_contexts( request, @@ -29744,6 +30081,7 @@ async def test_list_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29944,10 +30282,14 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateContextRequest.pb( metadata_service.UpdateContextRequest() ) @@ -29971,6 +30313,7 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_context.Context() + post_with_metadata.return_value = gca_context.Context(), metadata await client.update_context( request, @@ -29982,6 +30325,7 @@ async def test_update_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30082,10 +30426,14 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_context" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_context_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_context" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteContextRequest.pb( metadata_service.DeleteContextRequest() ) @@ -30109,6 +30457,7 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_context( request, @@ -30120,6 +30469,7 @@ async def test_delete_context_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30220,10 +30570,14 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_contexts" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeContextsRequest.pb( metadata_service.PurgeContextsRequest() ) @@ -30247,6 +30601,7 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_contexts( request, @@ -30258,6 +30613,7 @@ async def test_purge_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30368,11 +30724,15 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_add_context_artifacts_and_executions", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_context_artifacts_and_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_context_artifacts_and_executions", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextArtifactsAndExecutionsRequest.pb( metadata_service.AddContextArtifactsAndExecutionsRequest() ) @@ -30400,6 +30760,10 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextArtifactsAndExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextArtifactsAndExecutionsResponse(), + metadata, + ) await client.add_context_artifacts_and_executions( request, @@ -30411,6 +30775,7 @@ async def test_add_context_artifacts_and_executions_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30512,10 +30877,14 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_add_context_children" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddContextChildrenRequest.pb( metadata_service.AddContextChildrenRequest() ) @@ -30541,6 +30910,10 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.AddContextChildrenResponse(), + metadata, + ) await client.add_context_children( request, @@ -30552,6 +30925,7 @@ async def test_add_context_children_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30653,10 +31027,14 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_remove_context_children" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_remove_context_children_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_remove_context_children" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.RemoveContextChildrenRequest.pb( metadata_service.RemoveContextChildrenRequest() ) @@ -30682,6 +31060,10 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = metadata_service.RemoveContextChildrenResponse() + post_with_metadata.return_value = ( + metadata_service.RemoveContextChildrenResponse(), + metadata, + ) await client.remove_context_children( request, @@ -30693,6 +31075,7 @@ async def test_remove_context_children_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30797,11 +31180,15 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_context_lineage_subgraph", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_context_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_context_lineage_subgraph", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryContextLineageSubgraphRequest.pb( metadata_service.QueryContextLineageSubgraphRequest() ) @@ -30827,6 +31214,7 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_context_lineage_subgraph( request, @@ -30838,6 +31226,7 @@ async def test_query_context_lineage_subgraph_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31034,10 +31423,14 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateExecutionRequest.pb( metadata_service.CreateExecutionRequest() ) @@ -31061,6 +31454,7 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata await client.create_execution( request, @@ -31072,6 +31466,7 @@ async def test_create_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31188,10 +31583,14 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetExecutionRequest.pb( metadata_service.GetExecutionRequest() ) @@ -31215,6 +31614,7 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = execution.Execution() + post_with_metadata.return_value = execution.Execution(), metadata await client.get_execution( request, @@ -31226,6 +31626,7 @@ async def test_get_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31330,10 +31731,14 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_executions" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListExecutionsRequest.pb( metadata_service.ListExecutionsRequest() ) @@ -31359,6 +31764,10 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.ListExecutionsResponse() + post_with_metadata.return_value = ( + metadata_service.ListExecutionsResponse(), + metadata, + ) await client.list_executions( request, @@ -31370,6 +31779,7 @@ async def test_list_executions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31570,10 +31980,14 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_update_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_update_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_update_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.UpdateExecutionRequest.pb( metadata_service.UpdateExecutionRequest() ) @@ -31597,6 +32011,7 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_execution.Execution() + post_with_metadata.return_value = gca_execution.Execution(), metadata await client.update_execution( request, @@ -31608,6 +32023,7 @@ async def test_update_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31708,10 +32124,14 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_delete_execution" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_delete_execution_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_delete_execution" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.DeleteExecutionRequest.pb( metadata_service.DeleteExecutionRequest() ) @@ -31735,6 +32155,7 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_execution( request, @@ -31746,6 +32167,7 @@ async def test_delete_execution_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31846,10 +32268,14 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_purge_executions" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_purge_executions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_purge_executions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.PurgeExecutionsRequest.pb( metadata_service.PurgeExecutionsRequest() ) @@ -31873,6 +32299,7 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.purge_executions( request, @@ -31884,6 +32311,7 @@ async def test_purge_executions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31985,10 +32413,14 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_add_execution_events" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_add_execution_events_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_add_execution_events" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.AddExecutionEventsRequest.pb( metadata_service.AddExecutionEventsRequest() ) @@ -32014,6 +32446,10 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_service.AddExecutionEventsResponse() + post_with_metadata.return_value = ( + metadata_service.AddExecutionEventsResponse(), + metadata, + ) await client.add_execution_events( request, @@ -32025,6 +32461,7 @@ async def test_add_execution_events_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32131,11 +32568,15 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_execution_inputs_and_outputs", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_execution_inputs_and_outputs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_execution_inputs_and_outputs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryExecutionInputsAndOutputsRequest.pb( metadata_service.QueryExecutionInputsAndOutputsRequest() ) @@ -32161,6 +32602,7 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_execution_inputs_and_outputs( request, @@ -32172,6 +32614,7 @@ async def test_query_execution_inputs_and_outputs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32364,10 +32807,14 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_create_metadata_schema" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_create_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_create_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.CreateMetadataSchemaRequest.pb( metadata_service.CreateMetadataSchemaRequest() ) @@ -32393,6 +32840,7 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_metadata_schema.MetadataSchema() + post_with_metadata.return_value = gca_metadata_schema.MetadataSchema(), metadata await client.create_metadata_schema( request, @@ -32404,6 +32852,7 @@ async def test_create_metadata_schema_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32519,10 +32968,14 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_get_metadata_schema" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_get_metadata_schema_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_get_metadata_schema" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.GetMetadataSchemaRequest.pb( metadata_service.GetMetadataSchemaRequest() ) @@ -32548,6 +33001,7 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = metadata_schema.MetadataSchema() + post_with_metadata.return_value = metadata_schema.MetadataSchema(), metadata await client.get_metadata_schema( request, @@ -32559,6 +33013,7 @@ async def test_get_metadata_schema_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32663,10 +33118,14 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "post_list_metadata_schemas" ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_list_metadata_schemas_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_list_metadata_schemas" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.ListMetadataSchemasRequest.pb( metadata_service.ListMetadataSchemasRequest() ) @@ -32692,6 +33151,10 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = metadata_service.ListMetadataSchemasResponse() + post_with_metadata.return_value = ( + metadata_service.ListMetadataSchemasResponse(), + metadata, + ) await client.list_metadata_schemas( request, @@ -32703,6 +33166,7 @@ async def test_list_metadata_schemas_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -32807,11 +33271,15 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( transports.AsyncMetadataServiceRestInterceptor, "post_query_artifact_lineage_subgraph", ) as post, mock.patch.object( + transports.AsyncMetadataServiceRestInterceptor, + "post_query_artifact_lineage_subgraph_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMetadataServiceRestInterceptor, "pre_query_artifact_lineage_subgraph", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = metadata_service.QueryArtifactLineageSubgraphRequest.pb( metadata_service.QueryArtifactLineageSubgraphRequest() ) @@ -32837,6 +33305,7 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = lineage_subgraph.LineageSubgraph() + post_with_metadata.return_value = lineage_subgraph.LineageSubgraph(), metadata await client.query_artifact_lineage_subgraph( request, @@ -32848,6 +33317,7 @@ async def test_query_artifact_lineage_subgraph_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_migration_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_migration_service.py index 80bf139e67..ce6b852459 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_migration_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_migration_service.py @@ -85,6 +85,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -343,6 +351,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = MigrationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = MigrationServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2834,10 +2885,14 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.MigrationServiceRestInterceptor, "post_search_migratable_resources" ) as post, mock.patch.object( + transports.MigrationServiceRestInterceptor, + "post_search_migratable_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MigrationServiceRestInterceptor, "pre_search_migratable_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.SearchMigratableResourcesRequest.pb( migration_service.SearchMigratableResourcesRequest() ) @@ -2863,6 +2918,10 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = migration_service.SearchMigratableResourcesResponse() + post_with_metadata.return_value = ( + migration_service.SearchMigratableResourcesResponse(), + metadata, + ) client.search_migratable_resources( request, @@ -2874,6 +2933,7 @@ def test_search_migratable_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_migrate_resources_rest_bad_request( @@ -2954,10 +3014,14 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): ), mock.patch.object( transports.MigrationServiceRestInterceptor, "post_batch_migrate_resources" ) as post, mock.patch.object( + transports.MigrationServiceRestInterceptor, + "post_batch_migrate_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.MigrationServiceRestInterceptor, "pre_batch_migrate_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.BatchMigrateResourcesRequest.pb( migration_service.BatchMigrateResourcesRequest() ) @@ -2981,6 +3045,7 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_migrate_resources( request, @@ -2992,6 +3057,7 @@ def test_batch_migrate_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3799,11 +3865,15 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc transports.AsyncMigrationServiceRestInterceptor, "post_search_migratable_resources", ) as post, mock.patch.object( + transports.AsyncMigrationServiceRestInterceptor, + "post_search_migratable_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "pre_search_migratable_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.SearchMigratableResourcesRequest.pb( migration_service.SearchMigratableResourcesRequest() ) @@ -3829,6 +3899,10 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = migration_service.SearchMigratableResourcesResponse() + post_with_metadata.return_value = ( + migration_service.SearchMigratableResourcesResponse(), + metadata, + ) await client.search_migratable_resources( request, @@ -3840,6 +3914,7 @@ async def test_search_migratable_resources_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3936,10 +4011,14 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "post_batch_migrate_resources" ) as post, mock.patch.object( + transports.AsyncMigrationServiceRestInterceptor, + "post_batch_migrate_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncMigrationServiceRestInterceptor, "pre_batch_migrate_resources" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = migration_service.BatchMigrateResourcesRequest.pb( migration_service.BatchMigrateResourcesRequest() ) @@ -3963,6 +4042,7 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_migrate_resources( request, @@ -3974,6 +4054,7 @@ async def test_batch_migrate_resources_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_model_garden_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_model_garden_service.py index 2a7ea4e7f5..303636c3dd 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_model_garden_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_model_garden_service.py @@ -89,6 +89,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -359,6 +367,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ModelGardenServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ModelGardenServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3337,10 +3388,14 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "post_get_publisher_model" ) as post, mock.patch.object( + transports.ModelGardenServiceRestInterceptor, + "post_get_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "pre_get_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.GetPublisherModelRequest.pb( model_garden_service.GetPublisherModelRequest() ) @@ -3366,6 +3421,7 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = publisher_model.PublisherModel() + post_with_metadata.return_value = publisher_model.PublisherModel(), metadata client.get_publisher_model( request, @@ -3377,6 +3433,7 @@ def test_get_publisher_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_publisher_models_rest_bad_request( @@ -3461,10 +3518,14 @@ def test_list_publisher_models_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "post_list_publisher_models" ) as post, mock.patch.object( + transports.ModelGardenServiceRestInterceptor, + "post_list_publisher_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "pre_list_publisher_models" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.ListPublisherModelsRequest.pb( model_garden_service.ListPublisherModelsRequest() ) @@ -3490,6 +3551,10 @@ def test_list_publisher_models_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_garden_service.ListPublisherModelsResponse() + post_with_metadata.return_value = ( + model_garden_service.ListPublisherModelsResponse(), + metadata, + ) client.list_publisher_models( request, @@ -3501,6 +3566,7 @@ def test_list_publisher_models_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_deploy_publisher_model_rest_bad_request( @@ -3581,10 +3647,14 @@ def test_deploy_publisher_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelGardenServiceRestInterceptor, "post_deploy_publisher_model" ) as post, mock.patch.object( + transports.ModelGardenServiceRestInterceptor, + "post_deploy_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelGardenServiceRestInterceptor, "pre_deploy_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.DeployPublisherModelRequest.pb( model_garden_service.DeployPublisherModelRequest() ) @@ -3608,6 +3678,7 @@ def test_deploy_publisher_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.deploy_publisher_model( request, @@ -3619,6 +3690,7 @@ def test_deploy_publisher_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -4465,10 +4537,14 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "post_get_publisher_model" ) as post, mock.patch.object( + transports.AsyncModelGardenServiceRestInterceptor, + "post_get_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "pre_get_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.GetPublisherModelRequest.pb( model_garden_service.GetPublisherModelRequest() ) @@ -4494,6 +4570,7 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = publisher_model.PublisherModel() + post_with_metadata.return_value = publisher_model.PublisherModel(), metadata await client.get_publisher_model( request, @@ -4505,6 +4582,7 @@ async def test_get_publisher_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4605,10 +4683,14 @@ async def test_list_publisher_models_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "post_list_publisher_models" ) as post, mock.patch.object( + transports.AsyncModelGardenServiceRestInterceptor, + "post_list_publisher_models_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "pre_list_publisher_models" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.ListPublisherModelsRequest.pb( model_garden_service.ListPublisherModelsRequest() ) @@ -4634,6 +4716,10 @@ async def test_list_publisher_models_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = model_garden_service.ListPublisherModelsResponse() + post_with_metadata.return_value = ( + model_garden_service.ListPublisherModelsResponse(), + metadata, + ) await client.list_publisher_models( request, @@ -4645,6 +4731,7 @@ async def test_list_publisher_models_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4741,10 +4828,14 @@ async def test_deploy_publisher_model_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "post_deploy_publisher_model" ) as post, mock.patch.object( + transports.AsyncModelGardenServiceRestInterceptor, + "post_deploy_publisher_model_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelGardenServiceRestInterceptor, "pre_deploy_publisher_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_garden_service.DeployPublisherModelRequest.pb( model_garden_service.DeployPublisherModelRequest() ) @@ -4768,6 +4859,7 @@ async def test_deploy_publisher_model_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.deploy_publisher_model( request, @@ -4779,6 +4871,7 @@ async def test_deploy_publisher_model_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_model_monitoring_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_model_monitoring_service.py index e5667505be..79d2bf6625 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_model_monitoring_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_model_monitoring_service.py @@ -108,6 +108,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -382,6 +390,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ModelMonitoringServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ModelMonitoringServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -9367,10 +9418,14 @@ def test_create_model_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "post_create_model_monitor" ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_create_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_create_model_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.CreateModelMonitorRequest.pb( model_monitoring_service.CreateModelMonitorRequest() ) @@ -9394,6 +9449,7 @@ def test_create_model_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_model_monitor( request, @@ -9405,6 +9461,7 @@ def test_create_model_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_model_monitor_rest_bad_request( @@ -9694,10 +9751,14 @@ def test_update_model_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "post_update_model_monitor" ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_update_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_update_model_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.UpdateModelMonitorRequest.pb( model_monitoring_service.UpdateModelMonitorRequest() ) @@ -9721,6 +9782,7 @@ def test_update_model_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_model_monitor( request, @@ -9732,6 +9794,7 @@ def test_update_model_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_monitor_rest_bad_request( @@ -9822,10 +9885,14 @@ def test_get_model_monitor_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "post_get_model_monitor" ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_get_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_get_model_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.GetModelMonitorRequest.pb( model_monitoring_service.GetModelMonitorRequest() ) @@ -9849,6 +9916,7 @@ def test_get_model_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitor.ModelMonitor() + post_with_metadata.return_value = model_monitor.ModelMonitor(), metadata client.get_model_monitor( request, @@ -9860,6 +9928,7 @@ def test_get_model_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_monitors_rest_bad_request( @@ -9946,10 +10015,14 @@ def test_list_model_monitors_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "post_list_model_monitors" ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_list_model_monitors_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_list_model_monitors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.ListModelMonitorsRequest.pb( model_monitoring_service.ListModelMonitorsRequest() ) @@ -9975,6 +10048,10 @@ def test_list_model_monitors_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitoring_service.ListModelMonitorsResponse() + post_with_metadata.return_value = ( + model_monitoring_service.ListModelMonitorsResponse(), + metadata, + ) client.list_model_monitors( request, @@ -9986,6 +10063,7 @@ def test_list_model_monitors_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_monitor_rest_bad_request( @@ -10066,10 +10144,14 @@ def test_delete_model_monitor_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "post_delete_model_monitor" ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_delete_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_delete_model_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.DeleteModelMonitorRequest.pb( model_monitoring_service.DeleteModelMonitorRequest() ) @@ -10093,6 +10175,7 @@ def test_delete_model_monitor_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_monitor( request, @@ -10104,6 +10187,7 @@ def test_delete_model_monitor_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_model_monitoring_job_rest_bad_request( @@ -10412,11 +10496,15 @@ def test_create_model_monitoring_job_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_create_model_monitoring_job", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_create_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_create_model_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.CreateModelMonitoringJobRequest.pb( model_monitoring_service.CreateModelMonitoringJobRequest() ) @@ -10442,6 +10530,10 @@ def test_create_model_monitoring_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model_monitoring_job.ModelMonitoringJob() + post_with_metadata.return_value = ( + gca_model_monitoring_job.ModelMonitoringJob(), + metadata, + ) client.create_model_monitoring_job( request, @@ -10453,6 +10545,7 @@ def test_create_model_monitoring_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_monitoring_job_rest_bad_request( @@ -10548,10 +10641,14 @@ def test_get_model_monitoring_job_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_get_model_monitoring_job", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_get_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_get_model_monitoring_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.GetModelMonitoringJobRequest.pb( model_monitoring_service.GetModelMonitoringJobRequest() ) @@ -10577,6 +10674,10 @@ def test_get_model_monitoring_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitoring_job.ModelMonitoringJob() + post_with_metadata.return_value = ( + model_monitoring_job.ModelMonitoringJob(), + metadata, + ) client.get_model_monitoring_job( request, @@ -10588,6 +10689,7 @@ def test_get_model_monitoring_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_monitoring_jobs_rest_bad_request( @@ -10679,11 +10781,15 @@ def test_list_model_monitoring_jobs_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_list_model_monitoring_jobs", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_list_model_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_list_model_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.ListModelMonitoringJobsRequest.pb( model_monitoring_service.ListModelMonitoringJobsRequest() ) @@ -10709,6 +10815,10 @@ def test_list_model_monitoring_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitoring_service.ListModelMonitoringJobsResponse() + post_with_metadata.return_value = ( + model_monitoring_service.ListModelMonitoringJobsResponse(), + metadata, + ) client.list_model_monitoring_jobs( request, @@ -10720,6 +10830,7 @@ def test_list_model_monitoring_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_monitoring_job_rest_bad_request( @@ -10805,11 +10916,15 @@ def test_delete_model_monitoring_job_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_delete_model_monitoring_job", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_delete_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_delete_model_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.DeleteModelMonitoringJobRequest.pb( model_monitoring_service.DeleteModelMonitoringJobRequest() ) @@ -10833,6 +10948,7 @@ def test_delete_model_monitoring_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_monitoring_job( request, @@ -10844,6 +10960,7 @@ def test_delete_model_monitoring_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_model_monitoring_stats_rest_bad_request( @@ -10935,11 +11052,15 @@ def test_search_model_monitoring_stats_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_search_model_monitoring_stats", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_search_model_monitoring_stats_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_search_model_monitoring_stats", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.SearchModelMonitoringStatsRequest.pb( model_monitoring_service.SearchModelMonitoringStatsRequest() ) @@ -10969,6 +11090,10 @@ def test_search_model_monitoring_stats_rest_interceptors(null_interceptor): post.return_value = ( model_monitoring_service.SearchModelMonitoringStatsResponse() ) + post_with_metadata.return_value = ( + model_monitoring_service.SearchModelMonitoringStatsResponse(), + metadata, + ) client.search_model_monitoring_stats( request, @@ -10980,6 +11105,7 @@ def test_search_model_monitoring_stats_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_search_model_monitoring_alerts_rest_bad_request( @@ -11073,11 +11199,15 @@ def test_search_model_monitoring_alerts_rest_interceptors(null_interceptor): transports.ModelMonitoringServiceRestInterceptor, "post_search_model_monitoring_alerts", ) as post, mock.patch.object( + transports.ModelMonitoringServiceRestInterceptor, + "post_search_model_monitoring_alerts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelMonitoringServiceRestInterceptor, "pre_search_model_monitoring_alerts", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.SearchModelMonitoringAlertsRequest.pb( model_monitoring_service.SearchModelMonitoringAlertsRequest() ) @@ -11107,6 +11237,10 @@ def test_search_model_monitoring_alerts_rest_interceptors(null_interceptor): post.return_value = ( model_monitoring_service.SearchModelMonitoringAlertsResponse() ) + post_with_metadata.return_value = ( + model_monitoring_service.SearchModelMonitoringAlertsResponse(), + metadata, + ) client.search_model_monitoring_alerts( request, @@ -11118,6 +11252,7 @@ def test_search_model_monitoring_alerts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -12318,11 +12453,15 @@ async def test_create_model_monitor_rest_asyncio_interceptors(null_interceptor): transports.AsyncModelMonitoringServiceRestInterceptor, "post_create_model_monitor", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_create_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_create_model_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.CreateModelMonitorRequest.pb( model_monitoring_service.CreateModelMonitorRequest() ) @@ -12346,6 +12485,7 @@ async def test_create_model_monitor_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_model_monitor( request, @@ -12357,6 +12497,7 @@ async def test_create_model_monitor_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12663,11 +12804,15 @@ async def test_update_model_monitor_rest_asyncio_interceptors(null_interceptor): transports.AsyncModelMonitoringServiceRestInterceptor, "post_update_model_monitor", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_update_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_update_model_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.UpdateModelMonitorRequest.pb( model_monitoring_service.UpdateModelMonitorRequest() ) @@ -12691,6 +12836,7 @@ async def test_update_model_monitor_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_model_monitor( request, @@ -12702,6 +12848,7 @@ async def test_update_model_monitor_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12808,10 +12955,14 @@ async def test_get_model_monitor_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "post_get_model_monitor" ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_get_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_get_model_monitor" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.GetModelMonitorRequest.pb( model_monitoring_service.GetModelMonitorRequest() ) @@ -12835,6 +12986,7 @@ async def test_get_model_monitor_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitor.ModelMonitor() + post_with_metadata.return_value = model_monitor.ModelMonitor(), metadata await client.get_model_monitor( request, @@ -12846,6 +12998,7 @@ async def test_get_model_monitor_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12949,10 +13102,14 @@ async def test_list_model_monitors_rest_asyncio_interceptors(null_interceptor): transports.AsyncModelMonitoringServiceRestInterceptor, "post_list_model_monitors", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_list_model_monitors_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_list_model_monitors" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.ListModelMonitorsRequest.pb( model_monitoring_service.ListModelMonitorsRequest() ) @@ -12978,6 +13135,10 @@ async def test_list_model_monitors_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_monitoring_service.ListModelMonitorsResponse() + post_with_metadata.return_value = ( + model_monitoring_service.ListModelMonitorsResponse(), + metadata, + ) await client.list_model_monitors( request, @@ -12989,6 +13150,7 @@ async def test_list_model_monitors_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13086,11 +13248,15 @@ async def test_delete_model_monitor_rest_asyncio_interceptors(null_interceptor): transports.AsyncModelMonitoringServiceRestInterceptor, "post_delete_model_monitor", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_delete_model_monitor_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_delete_model_monitor", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.DeleteModelMonitorRequest.pb( model_monitoring_service.DeleteModelMonitorRequest() ) @@ -13114,6 +13280,7 @@ async def test_delete_model_monitor_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_monitor( request, @@ -13125,6 +13292,7 @@ async def test_delete_model_monitor_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13449,11 +13617,15 @@ async def test_create_model_monitoring_job_rest_asyncio_interceptors(null_interc transports.AsyncModelMonitoringServiceRestInterceptor, "post_create_model_monitoring_job", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_create_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_create_model_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.CreateModelMonitoringJobRequest.pb( model_monitoring_service.CreateModelMonitoringJobRequest() ) @@ -13479,6 +13651,10 @@ async def test_create_model_monitoring_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = gca_model_monitoring_job.ModelMonitoringJob() + post_with_metadata.return_value = ( + gca_model_monitoring_job.ModelMonitoringJob(), + metadata, + ) await client.create_model_monitoring_job( request, @@ -13490,6 +13666,7 @@ async def test_create_model_monitoring_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13601,11 +13778,15 @@ async def test_get_model_monitoring_job_rest_asyncio_interceptors(null_intercept transports.AsyncModelMonitoringServiceRestInterceptor, "post_get_model_monitoring_job", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_get_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_get_model_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.GetModelMonitoringJobRequest.pb( model_monitoring_service.GetModelMonitoringJobRequest() ) @@ -13631,6 +13812,10 @@ async def test_get_model_monitoring_job_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = model_monitoring_job.ModelMonitoringJob() + post_with_metadata.return_value = ( + model_monitoring_job.ModelMonitoringJob(), + metadata, + ) await client.get_model_monitoring_job( request, @@ -13642,6 +13827,7 @@ async def test_get_model_monitoring_job_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13749,11 +13935,15 @@ async def test_list_model_monitoring_jobs_rest_asyncio_interceptors(null_interce transports.AsyncModelMonitoringServiceRestInterceptor, "post_list_model_monitoring_jobs", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_list_model_monitoring_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_list_model_monitoring_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.ListModelMonitoringJobsRequest.pb( model_monitoring_service.ListModelMonitoringJobsRequest() ) @@ -13779,6 +13969,10 @@ async def test_list_model_monitoring_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = model_monitoring_service.ListModelMonitoringJobsResponse() + post_with_metadata.return_value = ( + model_monitoring_service.ListModelMonitoringJobsResponse(), + metadata, + ) await client.list_model_monitoring_jobs( request, @@ -13790,6 +13984,7 @@ async def test_list_model_monitoring_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13891,11 +14086,15 @@ async def test_delete_model_monitoring_job_rest_asyncio_interceptors(null_interc transports.AsyncModelMonitoringServiceRestInterceptor, "post_delete_model_monitoring_job", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_delete_model_monitoring_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_delete_model_monitoring_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.DeleteModelMonitoringJobRequest.pb( model_monitoring_service.DeleteModelMonitoringJobRequest() ) @@ -13919,6 +14118,7 @@ async def test_delete_model_monitoring_job_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_monitoring_job( request, @@ -13930,6 +14130,7 @@ async def test_delete_model_monitoring_job_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14039,11 +14240,15 @@ async def test_search_model_monitoring_stats_rest_asyncio_interceptors( transports.AsyncModelMonitoringServiceRestInterceptor, "post_search_model_monitoring_stats", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_search_model_monitoring_stats_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_search_model_monitoring_stats", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.SearchModelMonitoringStatsRequest.pb( model_monitoring_service.SearchModelMonitoringStatsRequest() ) @@ -14073,6 +14278,10 @@ async def test_search_model_monitoring_stats_rest_asyncio_interceptors( post.return_value = ( model_monitoring_service.SearchModelMonitoringStatsResponse() ) + post_with_metadata.return_value = ( + model_monitoring_service.SearchModelMonitoringStatsResponse(), + metadata, + ) await client.search_model_monitoring_stats( request, @@ -14084,6 +14293,7 @@ async def test_search_model_monitoring_stats_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14195,11 +14405,15 @@ async def test_search_model_monitoring_alerts_rest_asyncio_interceptors( transports.AsyncModelMonitoringServiceRestInterceptor, "post_search_model_monitoring_alerts", ) as post, mock.patch.object( + transports.AsyncModelMonitoringServiceRestInterceptor, + "post_search_model_monitoring_alerts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelMonitoringServiceRestInterceptor, "pre_search_model_monitoring_alerts", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_monitoring_service.SearchModelMonitoringAlertsRequest.pb( model_monitoring_service.SearchModelMonitoringAlertsRequest() ) @@ -14229,6 +14443,10 @@ async def test_search_model_monitoring_alerts_rest_asyncio_interceptors( post.return_value = ( model_monitoring_service.SearchModelMonitoringAlertsResponse() ) + post_with_metadata.return_value = ( + model_monitoring_service.SearchModelMonitoringAlertsResponse(), + metadata, + ) await client.search_model_monitoring_alerts( request, @@ -14240,6 +14458,7 @@ async def test_search_model_monitoring_alerts_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_model_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_model_service.py index 4ce5f5836a..2f1d5f4bf5 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_model_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_model_service.py @@ -103,6 +103,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -345,6 +353,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ModelServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ModelServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1449,6 +1500,7 @@ def test_get_model(request_type, transport: str = "grpc"): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -1478,6 +1530,7 @@ def test_get_model(request_type, transport: str = "grpc"): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -1623,6 +1676,7 @@ async def test_get_model_async( display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -1657,6 +1711,7 @@ async def test_get_model_async( assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -2892,11 +2947,11 @@ async def test_list_model_versions_async_pages(): @pytest.mark.parametrize( "request_type", [ - model_service.UpdateModelRequest, + model_service.ListModelVersionCheckpointsRequest, dict, ], ) -def test_update_model(request_type, transport: str = "grpc"): +def test_list_model_version_checkpoints(request_type, transport: str = "grpc"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -2907,63 +2962,27 @@ def test_update_model(request_type, transport: str = "grpc"): request = request_type() # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model( - name="name_value", - version_id="version_id_value", - version_aliases=["version_aliases_value"], - display_name="display_name_value", - description="description_value", - version_description="version_description_value", - metadata_schema_uri="metadata_schema_uri_value", - training_pipeline="training_pipeline_value", - artifact_uri="artifact_uri_value", - supported_deployment_resources_types=[ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ], - supported_input_storage_formats=["supported_input_storage_formats_value"], - supported_output_storage_formats=["supported_output_storage_formats_value"], - etag="etag_value", - metadata_artifact="metadata_artifact_value", - satisfies_pzs=True, - satisfies_pzi=True, + call.return_value = model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", ) - response = client.update_model(request) + response = client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) == 1 _, args, _ = call.mock_calls[0] - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() assert args[0] == request # Establish that the response is the type that we expect. - assert isinstance(response, gca_model.Model) - assert response.name == "name_value" - assert response.version_id == "version_id_value" - assert response.version_aliases == ["version_aliases_value"] - assert response.display_name == "display_name_value" - assert response.description == "description_value" - assert response.version_description == "version_description_value" - assert response.metadata_schema_uri == "metadata_schema_uri_value" - assert response.training_pipeline == "training_pipeline_value" - assert response.artifact_uri == "artifact_uri_value" - assert response.supported_deployment_resources_types == [ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ] - assert response.supported_input_storage_formats == [ - "supported_input_storage_formats_value" - ] - assert response.supported_output_storage_formats == [ - "supported_output_storage_formats_value" - ] - assert response.etag == "etag_value" - assert response.metadata_artifact == "metadata_artifact_value" - assert response.satisfies_pzs is True - assert response.satisfies_pzi is True + assert isinstance(response, pagers.ListModelVersionCheckpointsPager) + assert response.next_page_token == "next_page_token_value" -def test_update_model_non_empty_request_with_auto_populated_field(): +def test_list_model_version_checkpoints_non_empty_request_with_auto_populated_field(): # This test is a coverage failsafe to make sure that UUID4 fields are # automatically populated, according to AIP-4235, with non-empty requests. client = ModelServiceClient( @@ -2974,20 +2993,28 @@ def test_update_model_non_empty_request_with_auto_populated_field(): # Populate all string fields in the request which are not UUID4 # since we want to check that UUID4 are populated automatically # if they meet the requirements of AIP 4235. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest( + name="name_value", + page_token="page_token_value", + ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: call.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client.update_model(request=request) + client.list_model_version_checkpoints(request=request) call.assert_called() _, args, _ = call.mock_calls[0] - assert args[0] == model_service.UpdateModelRequest() + assert args[0] == model_service.ListModelVersionCheckpointsRequest( + name="name_value", + page_token="page_token_value", + ) -def test_update_model_use_cached_wrapped_rpc(): +def test_list_model_version_checkpoints_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -3001,21 +3028,26 @@ def test_update_model_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.update_model in client._transport._wrapped_methods + assert ( + client._transport.list_model_version_checkpoints + in client._transport._wrapped_methods + ) # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.update_model] = mock_rpc + client._transport._wrapped_methods[ + client._transport.list_model_version_checkpoints + ] = mock_rpc request = {} - client.update_model(request) + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.update_model(request) + client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 @@ -3023,7 +3055,7 @@ def test_update_model_use_cached_wrapped_rpc(): @pytest.mark.asyncio -async def test_update_model_async_use_cached_wrapped_rpc( +async def test_list_model_version_checkpoints_async_use_cached_wrapped_rpc( transport: str = "grpc_asyncio", ): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, @@ -3040,7 +3072,7 @@ async def test_update_model_async_use_cached_wrapped_rpc( # Ensure method has been cached assert ( - client._client._transport.update_model + client._client._transport.list_model_version_checkpoints in client._client._transport._wrapped_methods ) @@ -3048,16 +3080,16 @@ async def test_update_model_async_use_cached_wrapped_rpc( mock_rpc = mock.AsyncMock() mock_rpc.return_value = mock.Mock() client._client._transport._wrapped_methods[ - client._client._transport.update_model + client._client._transport.list_model_version_checkpoints ] = mock_rpc request = {} - await client.update_model(request) + await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - await client.update_model(request) + await client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 @@ -3065,8 +3097,9 @@ async def test_update_model_async_use_cached_wrapped_rpc( @pytest.mark.asyncio -async def test_update_model_async( - transport: str = "grpc_asyncio", request_type=model_service.UpdateModelRequest +async def test_list_model_version_checkpoints_async( + transport: str = "grpc_asyncio", + request_type=model_service.ListModelVersionCheckpointsRequest, ): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), @@ -3078,88 +3111,50 @@ async def test_update_model_async( request = request_type() # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( - gca_model.Model( - name="name_value", - version_id="version_id_value", - version_aliases=["version_aliases_value"], - display_name="display_name_value", - description="description_value", - version_description="version_description_value", - metadata_schema_uri="metadata_schema_uri_value", - training_pipeline="training_pipeline_value", - artifact_uri="artifact_uri_value", - supported_deployment_resources_types=[ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ], - supported_input_storage_formats=[ - "supported_input_storage_formats_value" - ], - supported_output_storage_formats=[ - "supported_output_storage_formats_value" - ], - etag="etag_value", - metadata_artifact="metadata_artifact_value", - satisfies_pzs=True, - satisfies_pzi=True, + model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", ) ) - response = await client.update_model(request) + response = await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) _, args, _ = call.mock_calls[0] - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() assert args[0] == request # Establish that the response is the type that we expect. - assert isinstance(response, gca_model.Model) - assert response.name == "name_value" - assert response.version_id == "version_id_value" - assert response.version_aliases == ["version_aliases_value"] - assert response.display_name == "display_name_value" - assert response.description == "description_value" - assert response.version_description == "version_description_value" - assert response.metadata_schema_uri == "metadata_schema_uri_value" - assert response.training_pipeline == "training_pipeline_value" - assert response.artifact_uri == "artifact_uri_value" - assert response.supported_deployment_resources_types == [ - gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES - ] - assert response.supported_input_storage_formats == [ - "supported_input_storage_formats_value" - ] - assert response.supported_output_storage_formats == [ - "supported_output_storage_formats_value" - ] - assert response.etag == "etag_value" - assert response.metadata_artifact == "metadata_artifact_value" - assert response.satisfies_pzs is True - assert response.satisfies_pzi is True + assert isinstance(response, pagers.ListModelVersionCheckpointsAsyncPager) + assert response.next_page_token == "next_page_token_value" @pytest.mark.asyncio -async def test_update_model_async_from_dict(): - await test_update_model_async(request_type=dict) +async def test_list_model_version_checkpoints_async_from_dict(): + await test_list_model_version_checkpoints_async(request_type=dict) -def test_update_model_field_headers(): +def test_list_model_version_checkpoints_field_headers(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) # Any value that is part of the HTTP/1.1 URI should be sent as # a field header. Set these to a non-empty value. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() - request.model.name = "name_value" + request.name = "name_value" # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: - call.return_value = gca_model.Model() - client.update_model(request) + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = model_service.ListModelVersionCheckpointsResponse() + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) == 1 @@ -3170,26 +3165,30 @@ def test_update_model_field_headers(): _, _, kw = call.mock_calls[0] assert ( "x-goog-request-params", - "model.name=name_value", + "name=name_value", ) in kw["metadata"] @pytest.mark.asyncio -async def test_update_model_field_headers_async(): +async def test_list_model_version_checkpoints_field_headers_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) # Any value that is part of the HTTP/1.1 URI should be sent as # a field header. Set these to a non-empty value. - request = model_service.UpdateModelRequest() + request = model_service.ListModelVersionCheckpointsRequest() - request.model.name = "name_value" + request.name = "name_value" # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: - call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) - await client.update_model(request) + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse() + ) + await client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert len(call.mock_calls) @@ -3200,39 +3199,37 @@ async def test_update_model_field_headers_async(): _, _, kw = call.mock_calls[0] assert ( "x-goog-request-params", - "model.name=name_value", + "name=name_value", ) in kw["metadata"] -def test_update_model_flattened(): +def test_list_model_version_checkpoints_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model() + call.return_value = model_service.ListModelVersionCheckpointsResponse() # Call the method with a truthy value for each flattened field, # using the keyword arguments to the method. - client.update_model( - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + client.list_model_version_checkpoints( + name="name_value", ) # Establish that the underlying call was made with the expected # request object values. assert len(call.mock_calls) == 1 _, args, _ = call.mock_calls[0] - arg = args[0].model - mock_val = gca_model.Model(name="name_value") - assert arg == mock_val - arg = args[0].update_mask - mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + arg = args[0].name + mock_val = "name_value" assert arg == mock_val -def test_update_model_flattened_error(): +def test_list_model_version_checkpoints_flattened_error(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), ) @@ -3240,46 +3237,45 @@ def test_update_model_flattened_error(): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.update_model( - model_service.UpdateModelRequest(), - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), + name="name_value", ) @pytest.mark.asyncio -async def test_update_model_flattened_async(): +async def test_list_model_version_checkpoints_flattened_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) # Mock the actual call within the gRPC stub, and fake the request. - with mock.patch.object(type(client.transport.update_model), "__call__") as call: + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: # Designate an appropriate return value for the call. - call.return_value = gca_model.Model() + call.return_value = model_service.ListModelVersionCheckpointsResponse() - call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse() + ) # Call the method with a truthy value for each flattened field, # using the keyword arguments to the method. - response = await client.update_model( - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + response = await client.list_model_version_checkpoints( + name="name_value", ) # Establish that the underlying call was made with the expected # request object values. assert len(call.mock_calls) _, args, _ = call.mock_calls[0] - arg = args[0].model - mock_val = gca_model.Model(name="name_value") - assert arg == mock_val - arg = args[0].update_mask - mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + arg = args[0].name + mock_val = "name_value" assert arg == mock_val @pytest.mark.asyncio -async def test_update_model_flattened_error_async(): +async def test_list_model_version_checkpoints_flattened_error_async(): client = ModelServiceAsyncClient( credentials=async_anonymous_credentials(), ) @@ -3287,41 +3283,655 @@ async def test_update_model_flattened_error_async(): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - await client.update_model( - model_service.UpdateModelRequest(), - model=gca_model.Model(name="name_value"), - update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + await client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), + name="name_value", ) -@pytest.mark.parametrize( - "request_type", - [ - model_service.UpdateExplanationDatasetRequest, - dict, - ], -) -def test_update_explanation_dataset(request_type, transport: str = "grpc"): +def test_list_model_version_checkpoints_pager(transport_name: str = "grpc"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), - transport=transport, + transport=transport_name, ) - # Everything is optional in proto3 as far as the runtime is concerned, - # and we are mocking out the actual API, so just send an empty request. - request = request_type() - # Mock the actual call within the gRPC stub, and fake the request. with mock.patch.object( - type(client.transport.update_explanation_dataset), "__call__" + type(client.transport.list_model_version_checkpoints), "__call__" ) as call: - # Designate an appropriate return value for the call. - call.return_value = operations_pb2.Operation(name="operations/spam") - response = client.update_explanation_dataset(request) + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) - # Establish that the underlying gRPC stub method was called. - assert len(call.mock_calls) == 1 - _, args, _ = call.mock_calls[0] + expected_metadata = () + retry = retries.Retry() + timeout = 5 + expected_metadata = tuple(expected_metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", ""),)), + ) + pager = client.list_model_version_checkpoints( + request={}, retry=retry, timeout=timeout + ) + + assert pager._metadata == expected_metadata + assert pager._retry == retry + assert pager._timeout == timeout + + results = list(pager) + assert len(results) == 6 + assert all(isinstance(i, model_service.ModelVersionCheckpoint) for i in results) + + +def test_list_model_version_checkpoints_pages(transport_name: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport_name, + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + pages = list(client.list_model_version_checkpoints(request={}).pages) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_async_pager(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), + "__call__", + new_callable=mock.AsyncMock, + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + async_pager = await client.list_model_version_checkpoints( + request={}, + ) + assert async_pager.next_page_token == "abc" + responses = [] + async for response in async_pager: # pragma: no branch + responses.append(response) + + assert len(responses) == 6 + assert all( + isinstance(i, model_service.ModelVersionCheckpoint) for i in responses + ) + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_async_pages(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), + "__call__", + new_callable=mock.AsyncMock, + ) as call: + # Set the response to a series of pages. + call.side_effect = ( + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + next_page_token="abc", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], + next_page_token="def", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + ], + next_page_token="ghi", + ), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + ], + ), + RuntimeError, + ) + pages = [] + # Workaround issue in python 3.9 related to code coverage by adding `# pragma: no branch` + # See https://github.com/googleapis/gapic-generator-python/pull/1174#issuecomment-1025132372 + async for page_ in ( # pragma: no branch + await client.list_model_version_checkpoints(request={}) + ).pages: + pages.append(page_) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.UpdateModelRequest, + dict, + ], +) +def test_update_model(request_type, transport: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model( + name="name_value", + version_id="version_id_value", + version_aliases=["version_aliases_value"], + display_name="display_name_value", + description="description_value", + version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", + metadata_schema_uri="metadata_schema_uri_value", + training_pipeline="training_pipeline_value", + artifact_uri="artifact_uri_value", + supported_deployment_resources_types=[ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ], + supported_input_storage_formats=["supported_input_storage_formats_value"], + supported_output_storage_formats=["supported_output_storage_formats_value"], + etag="etag_value", + metadata_artifact="metadata_artifact_value", + satisfies_pzs=True, + satisfies_pzi=True, + ) + response = client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = model_service.UpdateModelRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, gca_model.Model) + assert response.name == "name_value" + assert response.version_id == "version_id_value" + assert response.version_aliases == ["version_aliases_value"] + assert response.display_name == "display_name_value" + assert response.description == "description_value" + assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" + assert response.metadata_schema_uri == "metadata_schema_uri_value" + assert response.training_pipeline == "training_pipeline_value" + assert response.artifact_uri == "artifact_uri_value" + assert response.supported_deployment_resources_types == [ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ] + assert response.supported_input_storage_formats == [ + "supported_input_storage_formats_value" + ] + assert response.supported_output_storage_formats == [ + "supported_output_storage_formats_value" + ] + assert response.etag == "etag_value" + assert response.metadata_artifact == "metadata_artifact_value" + assert response.satisfies_pzs is True + assert response.satisfies_pzi is True + + +def test_update_model_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are + # automatically populated, according to AIP-4235, with non-empty requests. + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Populate all string fields in the request which are not UUID4 + # since we want to check that UUID4 are populated automatically + # if they meet the requirements of AIP 4235. + request = model_service.UpdateModelRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client.update_model(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == model_service.UpdateModelRequest() + + +def test_update_model_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.update_model in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[client._transport.update_model] = mock_rpc + request = {} + client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + client.update_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_update_model_async_use_cached_wrapped_rpc( + transport: str = "grpc_asyncio", +): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method_async.wrap_method") as wrapper_fn: + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert ( + client._client._transport.update_model + in client._client._transport._wrapped_methods + ) + + # Replace cached wrapped function with mock + mock_rpc = mock.AsyncMock() + mock_rpc.return_value = mock.Mock() + client._client._transport._wrapped_methods[ + client._client._transport.update_model + ] = mock_rpc + + request = {} + await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + await client.update_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +@pytest.mark.asyncio +async def test_update_model_async( + transport: str = "grpc_asyncio", request_type=model_service.UpdateModelRequest +): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + gca_model.Model( + name="name_value", + version_id="version_id_value", + version_aliases=["version_aliases_value"], + display_name="display_name_value", + description="description_value", + version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", + metadata_schema_uri="metadata_schema_uri_value", + training_pipeline="training_pipeline_value", + artifact_uri="artifact_uri_value", + supported_deployment_resources_types=[ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ], + supported_input_storage_formats=[ + "supported_input_storage_formats_value" + ], + supported_output_storage_formats=[ + "supported_output_storage_formats_value" + ], + etag="etag_value", + metadata_artifact="metadata_artifact_value", + satisfies_pzs=True, + satisfies_pzi=True, + ) + ) + response = await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = model_service.UpdateModelRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, gca_model.Model) + assert response.name == "name_value" + assert response.version_id == "version_id_value" + assert response.version_aliases == ["version_aliases_value"] + assert response.display_name == "display_name_value" + assert response.description == "description_value" + assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" + assert response.metadata_schema_uri == "metadata_schema_uri_value" + assert response.training_pipeline == "training_pipeline_value" + assert response.artifact_uri == "artifact_uri_value" + assert response.supported_deployment_resources_types == [ + gca_model.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ] + assert response.supported_input_storage_formats == [ + "supported_input_storage_formats_value" + ] + assert response.supported_output_storage_formats == [ + "supported_output_storage_formats_value" + ] + assert response.etag == "etag_value" + assert response.metadata_artifact == "metadata_artifact_value" + assert response.satisfies_pzs is True + assert response.satisfies_pzi is True + + +@pytest.mark.asyncio +async def test_update_model_async_from_dict(): + await test_update_model_async(request_type=dict) + + +def test_update_model_field_headers(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = model_service.UpdateModelRequest() + + request.model.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value = gca_model.Model() + client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "model.name=name_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_update_model_field_headers_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = model_service.UpdateModelRequest() + + request.model.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + await client.update_model(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "model.name=name_value", + ) in kw["metadata"] + + +def test_update_model_flattened(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model() + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.update_model( + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + arg = args[0].model + mock_val = gca_model.Model(name="name_value") + assert arg == mock_val + arg = args[0].update_mask + mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + assert arg == mock_val + + +def test_update_model_flattened_error(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.update_model( + model_service.UpdateModelRequest(), + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + +@pytest.mark.asyncio +async def test_update_model_flattened_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.update_model), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = gca_model.Model() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall(gca_model.Model()) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.update_model( + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + arg = args[0].model + mock_val = gca_model.Model(name="name_value") + assert arg == mock_val + arg = args[0].update_mask + mock_val = field_mask_pb2.FieldMask(paths=["paths_value"]) + assert arg == mock_val + + +@pytest.mark.asyncio +async def test_update_model_flattened_error_async(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.update_model( + model_service.UpdateModelRequest(), + model=gca_model.Model(name="name_value"), + update_mask=field_mask_pb2.FieldMask(paths=["paths_value"]), + ) + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.UpdateExplanationDatasetRequest, + dict, + ], +) +def test_update_explanation_dataset(request_type, transport: str = "grpc"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.update_explanation_dataset), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/spam") + response = client.update_explanation_dataset(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] request = model_service.UpdateExplanationDatasetRequest() assert args[0] == request @@ -4343,6 +4953,7 @@ def test_merge_version_aliases(request_type, transport: str = "grpc"): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -4372,6 +4983,7 @@ def test_merge_version_aliases(request_type, transport: str = "grpc"): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -4529,6 +5141,7 @@ async def test_merge_version_aliases_async( display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -4563,6 +5176,7 @@ async def test_merge_version_aliases_async( assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -8381,20 +8995,202 @@ def test_upload_model_rest_use_cached_wrapped_rpc(): # subsequent calls should use the cached wrapper wrapper_fn.reset_mock() - client.upload_model(request) + client.upload_model(request) + + # Establish that a new wrapper was not created for this call + assert wrapper_fn.call_count == 0 + assert mock_rpc.call_count == 2 + + +def test_upload_model_rest_required_fields( + request_type=model_service.UploadModelRequest, +): + transport_class = transports.ModelServiceRestTransport + + request_init = {} + request_init["parent"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson(pb_request, use_integers_for_enums=False) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).upload_model._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["parent"] = "parent_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).upload_model._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "parent" in jsonified_request + assert jsonified_request["parent"] == "parent_value" + + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "post", + "query_params": pb_request, + } + transcode_result["body"] = pb_request + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + + response = client.upload_model(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_upload_model_rest_unset_required_fields(): + transport = transports.ModelServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.upload_model._get_unset_required_fields({}) + assert set(unset_fields) == ( + set(()) + & set( + ( + "parent", + "model", + ) + ) + ) + + +def test_upload_model_rest_flattened(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation(name="operations/spam") + + # get arguments that satisfy an http rule for this method + sample_request = {"parent": "projects/sample1/locations/sample2"} + + # get truthy value for each flattened field + mock_args = dict( + parent="parent_value", + model=gca_model.Model(name="name_value"), + ) + mock_args.update(sample_request) + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + + client.upload_model(**mock_args) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(req.mock_calls) == 1 + _, args, _ = req.mock_calls[0] + assert path_template.validate( + "%s/v1beta1/{parent=projects/*/locations/*}/models:upload" + % client.transport._host, + args[1], + ) + + +def test_upload_model_rest_flattened_error(transport: str = "rest"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.upload_model( + model_service.UploadModelRequest(), + parent="parent_value", + model=gca_model.Model(name="name_value"), + ) + + +def test_get_model_rest_use_cached_wrapped_rpc(): + # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, + # instead of constructing them on each call + with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Should wrap all calls on client creation + assert wrapper_fn.call_count > 0 + wrapper_fn.reset_mock() + + # Ensure method has been cached + assert client._transport.get_model in client._transport._wrapped_methods + + # Replace cached wrapped function with mock + mock_rpc = mock.Mock() + mock_rpc.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client._transport._wrapped_methods[client._transport.get_model] = mock_rpc + + request = {} + client.get_model(request) + + # Establish that the underlying gRPC stub method was called. + assert mock_rpc.call_count == 1 + + client.get_model(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_upload_model_rest_required_fields( - request_type=model_service.UploadModelRequest, -): +def test_get_model_rest_required_fields(request_type=model_service.GetModelRequest): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["parent"] = "" + request_init["name"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8405,21 +9201,21 @@ def test_upload_model_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).upload_model._get_unset_required_fields(jsonified_request) + ).get_model._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["parent"] = "parent_value" + jsonified_request["name"] = "name_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).upload_model._get_unset_required_fields(jsonified_request) + ).get_model._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "parent" in jsonified_request - assert jsonified_request["parent"] == "parent_value" + assert "name" in jsonified_request + assert jsonified_request["name"] == "name_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8428,7 +9224,7 @@ def test_upload_model_rest_required_fields( request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = operations_pb2.Operation(name="operations/spam") + return_value = model.Model() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8440,45 +9236,39 @@ def test_upload_model_rest_required_fields( pb_request = request_type.pb(request) transcode_result = { "uri": "v1/sample_method", - "method": "post", + "method": "get", "query_params": pb_request, } - transcode_result["body"] = pb_request transcode.return_value = transcode_result response_value = Response() response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model.Model.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.upload_model(request) + response = client.get_model(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_upload_model_rest_unset_required_fields(): +def test_get_model_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.upload_model._get_unset_required_fields({}) - assert set(unset_fields) == ( - set(()) - & set( - ( - "parent", - "model", - ) - ) - ) + unset_fields = transport.get_model._get_unset_required_fields({}) + assert set(unset_fields) == (set(()) & set(("name",))) -def test_upload_model_rest_flattened(): +def test_get_model_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8487,40 +9277,41 @@ def test_upload_model_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = operations_pb2.Operation(name="operations/spam") + return_value = model.Model() # get arguments that satisfy an http rule for this method - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} # get truthy value for each flattened field mock_args = dict( - parent="parent_value", - model=gca_model.Model(name="name_value"), + name="name_value", ) mock_args.update(sample_request) # Wrap the value into a proper Response obj response_value = Response() response_value.status_code = 200 + # Convert return value to protobuf type + return_value = model.Model.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.upload_model(**mock_args) + client.get_model(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1beta1/{parent=projects/*/locations/*}/models:upload" + "%s/v1beta1/{name=projects/*/locations/*/models/*}" % client.transport._host, args[1], ) -def test_upload_model_rest_flattened_error(transport: str = "rest"): +def test_get_model_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8529,14 +9320,13 @@ def test_upload_model_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.upload_model( - model_service.UploadModelRequest(), - parent="parent_value", - model=gca_model.Model(name="name_value"), + client.get_model( + model_service.GetModelRequest(), + name="name_value", ) -def test_get_model_rest_use_cached_wrapped_rpc(): +def test_list_models_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -8550,33 +9340,33 @@ def test_get_model_rest_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.get_model in client._transport._wrapped_methods + assert client._transport.list_models in client._transport._wrapped_methods # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.get_model] = mock_rpc + client._transport._wrapped_methods[client._transport.list_models] = mock_rpc request = {} - client.get_model(request) + client.list_models(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.get_model(request) + client.list_models(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_get_model_rest_required_fields(request_type=model_service.GetModelRequest): +def test_list_models_rest_required_fields(request_type=model_service.ListModelsRequest): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["name"] = "" + request_init["parent"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8587,21 +9377,30 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).get_model._get_unset_required_fields(jsonified_request) + ).list_models._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["name"] = "name_value" + jsonified_request["parent"] = "parent_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).get_model._get_unset_required_fields(jsonified_request) + ).list_models._get_unset_required_fields(jsonified_request) + # Check that path parameters and body parameters are not mixing in. + assert not set(unset_fields) - set( + ( + "filter", + "page_size", + "page_token", + "read_mask", + ) + ) jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "name" in jsonified_request - assert jsonified_request["name"] == "name_value" + assert "parent" in jsonified_request + assert jsonified_request["parent"] == "parent_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8610,7 +9409,7 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model.Model() + return_value = model_service.ListModelsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8631,30 +9430,40 @@ def test_get_model_rest_required_fields(request_type=model_service.GetModelReque response_value.status_code = 200 # Convert return value to protobuf type - return_value = model.Model.pb(return_value) + return_value = model_service.ListModelsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.get_model(request) + response = client.list_models(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_get_model_rest_unset_required_fields(): +def test_list_models_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.get_model._get_unset_required_fields({}) - assert set(unset_fields) == (set(()) & set(("name",))) + unset_fields = transport.list_models._get_unset_required_fields({}) + assert set(unset_fields) == ( + set( + ( + "filter", + "pageSize", + "pageToken", + "readMask", + ) + ) + & set(("parent",)) + ) -def test_get_model_rest_flattened(): +def test_list_models_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8663,14 +9472,14 @@ def test_get_model_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model.Model() + return_value = model_service.ListModelsResponse() # get arguments that satisfy an http rule for this method - sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} + sample_request = {"parent": "projects/sample1/locations/sample2"} # get truthy value for each flattened field mock_args = dict( - name="name_value", + parent="parent_value", ) mock_args.update(sample_request) @@ -8678,26 +9487,26 @@ def test_get_model_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model.Model.pb(return_value) + return_value = model_service.ListModelsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.get_model(**mock_args) + client.list_models(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1beta1/{name=projects/*/locations/*/models/*}" + "%s/v1beta1/{parent=projects/*/locations/*}/models" % client.transport._host, args[1], ) -def test_get_model_rest_flattened_error(transport: str = "rest"): +def test_list_models_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8706,13 +9515,74 @@ def test_get_model_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.get_model( - model_service.GetModelRequest(), - name="name_value", + client.list_models( + model_service.ListModelsRequest(), + parent="parent_value", ) -def test_list_models_rest_use_cached_wrapped_rpc(): +def test_list_models_rest_pager(transport: str = "rest"): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # TODO(kbandes): remove this mock unless there's a good reason for it. + # with mock.patch.object(path_template, 'transcode') as transcode: + # Set the response as a series of pages + response = ( + model_service.ListModelsResponse( + models=[ + model.Model(), + model.Model(), + model.Model(), + ], + next_page_token="abc", + ), + model_service.ListModelsResponse( + models=[], + next_page_token="def", + ), + model_service.ListModelsResponse( + models=[ + model.Model(), + ], + next_page_token="ghi", + ), + model_service.ListModelsResponse( + models=[ + model.Model(), + model.Model(), + ], + ), + ) + # Two responses for two calls + response = response + response + + # Wrap the values into proper Response objs + response = tuple(model_service.ListModelsResponse.to_json(x) for x in response) + return_values = tuple(Response() for i in response) + for return_val, response_val in zip(return_values, response): + return_val._content = response_val.encode("UTF-8") + return_val.status_code = 200 + req.side_effect = return_values + + sample_request = {"parent": "projects/sample1/locations/sample2"} + + pager = client.list_models(request=sample_request) + + results = list(pager) + assert len(results) == 6 + assert all(isinstance(i, model.Model) for i in results) + + pages = list(client.list_models(request=sample_request).pages) + for page_, token in zip(pages, ["abc", "def", "ghi", ""]): + assert page_.raw_page.next_page_token == token + + +def test_list_model_versions_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -8726,33 +9596,39 @@ def test_list_models_rest_use_cached_wrapped_rpc(): wrapper_fn.reset_mock() # Ensure method has been cached - assert client._transport.list_models in client._transport._wrapped_methods + assert ( + client._transport.list_model_versions in client._transport._wrapped_methods + ) # Replace cached wrapped function with mock mock_rpc = mock.Mock() mock_rpc.return_value.name = ( "foo" # operation_request.operation in compute client(s) expect a string. ) - client._transport._wrapped_methods[client._transport.list_models] = mock_rpc + client._transport._wrapped_methods[ + client._transport.list_model_versions + ] = mock_rpc request = {} - client.list_models(request) + client.list_model_versions(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.list_models(request) + client.list_model_versions(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_list_models_rest_required_fields(request_type=model_service.ListModelsRequest): +def test_list_model_versions_rest_required_fields( + request_type=model_service.ListModelVersionsRequest, +): transport_class = transports.ModelServiceRestTransport request_init = {} - request_init["parent"] = "" + request_init["name"] = "" request = request_type(**request_init) pb_request = request_type.pb(request) jsonified_request = json.loads( @@ -8763,20 +9639,21 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_models._get_unset_required_fields(jsonified_request) + ).list_model_versions._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present - jsonified_request["parent"] = "parent_value" + jsonified_request["name"] = "name_value" unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_models._get_unset_required_fields(jsonified_request) + ).list_model_versions._get_unset_required_fields(jsonified_request) # Check that path parameters and body parameters are not mixing in. assert not set(unset_fields) - set( ( "filter", + "order_by", "page_size", "page_token", "read_mask", @@ -8785,8 +9662,8 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR jsonified_request.update(unset_fields) # verify required fields with non-default values are left alone - assert "parent" in jsonified_request - assert jsonified_request["parent"] == "parent_value" + assert "name" in jsonified_request + assert jsonified_request["name"] == "name_value" client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), @@ -8795,7 +9672,7 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -8816,40 +9693,41 @@ def test_list_models_rest_required_fields(request_type=model_service.ListModelsR response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelsResponse.pb(return_value) + return_value = model_service.ListModelVersionsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_models(request) + response = client.list_model_versions(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_list_models_rest_unset_required_fields(): +def test_list_model_versions_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.list_models._get_unset_required_fields({}) + unset_fields = transport.list_model_versions._get_unset_required_fields({}) assert set(unset_fields) == ( set( ( "filter", + "orderBy", "pageSize", "pageToken", "readMask", ) ) - & set(("parent",)) + & set(("name",)) ) -def test_list_models_rest_flattened(): +def test_list_model_versions_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -8858,14 +9736,14 @@ def test_list_models_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse() # get arguments that satisfy an http rule for this method - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} # get truthy value for each flattened field mock_args = dict( - parent="parent_value", + name="name_value", ) mock_args.update(sample_request) @@ -8873,26 +9751,26 @@ def test_list_models_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelsResponse.pb(return_value) + return_value = model_service.ListModelVersionsResponse.pb(return_value) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_models(**mock_args) + client.list_model_versions(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1beta1/{parent=projects/*/locations/*}/models" + "%s/v1beta1/{name=projects/*/locations/*/models/*}:listVersions" % client.transport._host, args[1], ) -def test_list_models_rest_flattened_error(transport: str = "rest"): +def test_list_model_versions_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8901,13 +9779,13 @@ def test_list_models_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.list_models( - model_service.ListModelsRequest(), - parent="parent_value", + client.list_model_versions( + model_service.ListModelVersionsRequest(), + name="name_value", ) -def test_list_models_rest_pager(transport: str = "rest"): +def test_list_model_versions_rest_pager(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -8919,7 +9797,7 @@ def test_list_models_rest_pager(transport: str = "rest"): # with mock.patch.object(path_template, 'transcode') as transcode: # Set the response as a series of pages response = ( - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), model.Model(), @@ -8927,17 +9805,17 @@ def test_list_models_rest_pager(transport: str = "rest"): ], next_page_token="abc", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[], next_page_token="def", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), ], next_page_token="ghi", ), - model_service.ListModelsResponse( + model_service.ListModelVersionsResponse( models=[ model.Model(), model.Model(), @@ -8948,27 +9826,29 @@ def test_list_models_rest_pager(transport: str = "rest"): response = response + response # Wrap the values into proper Response objs - response = tuple(model_service.ListModelsResponse.to_json(x) for x in response) + response = tuple( + model_service.ListModelVersionsResponse.to_json(x) for x in response + ) return_values = tuple(Response() for i in response) for return_val, response_val in zip(return_values, response): return_val._content = response_val.encode("UTF-8") return_val.status_code = 200 req.side_effect = return_values - sample_request = {"parent": "projects/sample1/locations/sample2"} + sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} - pager = client.list_models(request=sample_request) + pager = client.list_model_versions(request=sample_request) results = list(pager) assert len(results) == 6 assert all(isinstance(i, model.Model) for i in results) - pages = list(client.list_models(request=sample_request).pages) + pages = list(client.list_model_versions(request=sample_request).pages) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token -def test_list_model_versions_rest_use_cached_wrapped_rpc(): +def test_list_model_version_checkpoints_rest_use_cached_wrapped_rpc(): # Clients should use _prep_wrapped_messages to create cached wrapped rpcs, # instead of constructing them on each call with mock.patch("google.api_core.gapic_v1.method.wrap_method") as wrapper_fn: @@ -8983,7 +9863,8 @@ def test_list_model_versions_rest_use_cached_wrapped_rpc(): # Ensure method has been cached assert ( - client._transport.list_model_versions in client._transport._wrapped_methods + client._transport.list_model_version_checkpoints + in client._transport._wrapped_methods ) # Replace cached wrapped function with mock @@ -8992,24 +9873,24 @@ def test_list_model_versions_rest_use_cached_wrapped_rpc(): "foo" # operation_request.operation in compute client(s) expect a string. ) client._transport._wrapped_methods[ - client._transport.list_model_versions + client._transport.list_model_version_checkpoints ] = mock_rpc request = {} - client.list_model_versions(request) + client.list_model_version_checkpoints(request) # Establish that the underlying gRPC stub method was called. assert mock_rpc.call_count == 1 - client.list_model_versions(request) + client.list_model_version_checkpoints(request) # Establish that a new wrapper was not created for this call assert wrapper_fn.call_count == 0 assert mock_rpc.call_count == 2 -def test_list_model_versions_rest_required_fields( - request_type=model_service.ListModelVersionsRequest, +def test_list_model_version_checkpoints_rest_required_fields( + request_type=model_service.ListModelVersionCheckpointsRequest, ): transport_class = transports.ModelServiceRestTransport @@ -9025,7 +9906,7 @@ def test_list_model_versions_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_model_versions._get_unset_required_fields(jsonified_request) + ).list_model_version_checkpoints._get_unset_required_fields(jsonified_request) jsonified_request.update(unset_fields) # verify required fields with default values are now present @@ -9034,15 +9915,12 @@ def test_list_model_versions_rest_required_fields( unset_fields = transport_class( credentials=ga_credentials.AnonymousCredentials() - ).list_model_versions._get_unset_required_fields(jsonified_request) + ).list_model_version_checkpoints._get_unset_required_fields(jsonified_request) # Check that path parameters and body parameters are not mixing in. assert not set(unset_fields) - set( ( - "filter", - "order_by", "page_size", "page_token", - "read_mask", ) ) jsonified_request.update(unset_fields) @@ -9058,7 +9936,7 @@ def test_list_model_versions_rest_required_fields( request = request_type(**request_init) # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse() # Mock the http request call within the method and fake a response. with mock.patch.object(Session, "request") as req: # We need to mock transcode() because providing default values @@ -9079,41 +9957,42 @@ def test_list_model_versions_rest_required_fields( response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_model_versions(request) + response = client.list_model_version_checkpoints(request) expected_params = [("$alt", "json;enum-encoding=int")] actual_params = req.call_args.kwargs["params"] assert expected_params == actual_params -def test_list_model_versions_rest_unset_required_fields(): +def test_list_model_version_checkpoints_rest_unset_required_fields(): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials ) - unset_fields = transport.list_model_versions._get_unset_required_fields({}) + unset_fields = transport.list_model_version_checkpoints._get_unset_required_fields( + {} + ) assert set(unset_fields) == ( set( ( - "filter", - "orderBy", "pageSize", "pageToken", - "readMask", ) ) & set(("name",)) ) -def test_list_model_versions_rest_flattened(): +def test_list_model_version_checkpoints_rest_flattened(): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest", @@ -9122,7 +10001,7 @@ def test_list_model_versions_rest_flattened(): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse() # get arguments that satisfy an http rule for this method sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} @@ -9137,26 +10016,28 @@ def test_list_model_versions_rest_flattened(): response_value = Response() response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value._content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_model_versions(**mock_args) + client.list_model_version_checkpoints(**mock_args) # Establish that the underlying call was made with the expected # request object values. assert len(req.mock_calls) == 1 _, args, _ = req.mock_calls[0] assert path_template.validate( - "%s/v1beta1/{name=projects/*/locations/*/models/*}:listVersions" + "%s/v1beta1/{name=projects/*/locations/*/models/*}:listCheckpoints" % client.transport._host, args[1], ) -def test_list_model_versions_rest_flattened_error(transport: str = "rest"): +def test_list_model_version_checkpoints_rest_flattened_error(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -9165,13 +10046,13 @@ def test_list_model_versions_rest_flattened_error(transport: str = "rest"): # Attempting to call a method with both a request object and flattened # fields is an error. with pytest.raises(ValueError): - client.list_model_versions( - model_service.ListModelVersionsRequest(), + client.list_model_version_checkpoints( + model_service.ListModelVersionCheckpointsRequest(), name="name_value", ) -def test_list_model_versions_rest_pager(transport: str = "rest"): +def test_list_model_version_checkpoints_rest_pager(transport: str = "rest"): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport=transport, @@ -9183,28 +10064,28 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): # with mock.patch.object(path_template, 'transcode') as transcode: # Set the response as a series of pages response = ( - model_service.ListModelVersionsResponse( - models=[ - model.Model(), - model.Model(), - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), ], next_page_token="abc", ), - model_service.ListModelVersionsResponse( - models=[], + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[], next_page_token="def", ), - model_service.ListModelVersionsResponse( - models=[ - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), ], next_page_token="ghi", ), - model_service.ListModelVersionsResponse( - models=[ - model.Model(), - model.Model(), + model_service.ListModelVersionCheckpointsResponse( + checkpoints=[ + model_service.ModelVersionCheckpoint(), + model_service.ModelVersionCheckpoint(), ], ), ) @@ -9213,7 +10094,8 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): # Wrap the values into proper Response objs response = tuple( - model_service.ListModelVersionsResponse.to_json(x) for x in response + model_service.ListModelVersionCheckpointsResponse.to_json(x) + for x in response ) return_values = tuple(Response() for i in response) for return_val, response_val in zip(return_values, response): @@ -9223,13 +10105,15 @@ def test_list_model_versions_rest_pager(transport: str = "rest"): sample_request = {"name": "projects/sample1/locations/sample2/models/sample3"} - pager = client.list_model_versions(request=sample_request) + pager = client.list_model_version_checkpoints(request=sample_request) results = list(pager) assert len(results) == 6 - assert all(isinstance(i, model.Model) for i in results) + assert all(isinstance(i, model_service.ModelVersionCheckpoint) for i in results) - pages = list(client.list_model_versions(request=sample_request).pages) + pages = list( + client.list_model_version_checkpoints(request=sample_request).pages + ) for page_, token in zip(pages, ["abc", "def", "ghi", ""]): assert page_.raw_page.next_page_token == token @@ -12256,6 +13140,29 @@ def test_list_model_versions_empty_call_grpc(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_list_model_version_checkpoints_empty_call_grpc(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + call.return_value = model_service.ListModelVersionCheckpointsResponse() + client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. def test_update_model_empty_call_grpc(): @@ -12629,6 +13536,7 @@ async def test_get_model_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -12713,6 +13621,35 @@ async def test_list_model_versions_empty_call_grpc_asyncio(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_empty_call_grpc_asyncio(): + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", + ) + ) + await client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @pytest.mark.asyncio @@ -12733,6 +13670,7 @@ async def test_update_model_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -12862,6 +13800,7 @@ async def test_merge_version_aliases_empty_call_grpc_asyncio(): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -13235,10 +14174,13 @@ def test_upload_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_upload_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_upload_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_upload_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UploadModelRequest.pb( model_service.UploadModelRequest() ) @@ -13262,6 +14204,7 @@ def test_upload_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.upload_model( request, @@ -13273,6 +14216,7 @@ def test_upload_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_rest_bad_request(request_type=model_service.GetModelRequest): @@ -13324,6 +14268,7 @@ def test_get_model_rest_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -13358,6 +14303,7 @@ def test_get_model_rest_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -13393,10 +14339,13 @@ def test_get_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_get_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelRequest.pb(model_service.GetModelRequest()) transcode.return_value = { "method": "post", @@ -13418,6 +14367,7 @@ def test_get_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata client.get_model( request, @@ -13429,6 +14379,7 @@ def test_get_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_models_rest_bad_request(request_type=model_service.ListModelsRequest): @@ -13487,15 +14438,144 @@ def test_list_models_rest_call_success(request_type): response_value.content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_models(request) + response = client.list_models(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, pagers.ListModelsPager) + assert response.next_page_token == "next_page_token_value" + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_list_models_rest_interceptors(null_interceptor): + transport = transports.ModelServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.ModelServiceRestInterceptor(), + ) + client = ModelServiceClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_list_models" + ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_list_models_with_metadata" + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_models" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelsRequest.pb( + model_service.ListModelsRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = model_service.ListModelsResponse.to_json( + model_service.ListModelsResponse() + ) + req.return_value.content = return_value + + request = model_service.ListModelsRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = model_service.ListModelsResponse() + post_with_metadata.return_value = model_service.ListModelsResponse(), metadata + + client.list_models( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() + + +def test_list_model_versions_rest_bad_request( + request_type=model_service.ListModelVersionsRequest, +): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + json_return_value = "" + response_value.json = mock.Mock(return_value={}) + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + client.list_model_versions(request) + + +@pytest.mark.parametrize( + "request_type", + [ + model_service.ListModelVersionsRequest, + dict, + ], +) +def test_list_model_versions_rest_call_success(request_type): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = model_service.ListModelVersionsResponse( + next_page_token="next_page_token_value", + ) + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model_service.ListModelVersionsResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + response_value.content = json_return_value.encode("UTF-8") + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = client.list_model_versions(request) # Establish that the response is the type that we expect. - assert isinstance(response, pagers.ListModelsPager) + assert isinstance(response, pagers.ListModelVersionsPager) assert response.next_page_token == "next_page_token_value" @pytest.mark.parametrize("null_interceptor", [True, False]) -def test_list_models_rest_interceptors(null_interceptor): +def test_list_model_versions_rest_interceptors(null_interceptor): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials(), interceptor=None @@ -13509,14 +14589,17 @@ def test_list_models_rest_interceptors(null_interceptor): ) as req, mock.patch.object( path_template, "transcode" ) as transcode, mock.patch.object( - transports.ModelServiceRestInterceptor, "post_list_models" + transports.ModelServiceRestInterceptor, "post_list_model_versions" ) as post, mock.patch.object( - transports.ModelServiceRestInterceptor, "pre_list_models" + transports.ModelServiceRestInterceptor, "post_list_model_versions_with_metadata" + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_model_versions" ) as pre: pre.assert_not_called() post.assert_not_called() - pb_message = model_service.ListModelsRequest.pb( - model_service.ListModelsRequest() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionsRequest.pb( + model_service.ListModelVersionsRequest() ) transcode.return_value = { "method": "post", @@ -13528,20 +14611,24 @@ def test_list_models_rest_interceptors(null_interceptor): req.return_value = mock.Mock() req.return_value.status_code = 200 req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - return_value = model_service.ListModelsResponse.to_json( - model_service.ListModelsResponse() + return_value = model_service.ListModelVersionsResponse.to_json( + model_service.ListModelVersionsResponse() ) req.return_value.content = return_value - request = model_service.ListModelsRequest() + request = model_service.ListModelVersionsRequest() metadata = [ ("key", "val"), ("cephalopod", "squid"), ] pre.return_value = request, metadata - post.return_value = model_service.ListModelsResponse() + post.return_value = model_service.ListModelVersionsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionsResponse(), + metadata, + ) - client.list_models( + client.list_model_versions( request, metadata=[ ("key", "val"), @@ -13551,10 +14638,11 @@ def test_list_models_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() -def test_list_model_versions_rest_bad_request( - request_type=model_service.ListModelVersionsRequest, +def test_list_model_version_checkpoints_rest_bad_request( + request_type=model_service.ListModelVersionCheckpointsRequest, ): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" @@ -13575,17 +14663,17 @@ def test_list_model_versions_rest_bad_request( response_value.request = mock.Mock() req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - client.list_model_versions(request) + client.list_model_version_checkpoints(request) @pytest.mark.parametrize( "request_type", [ - model_service.ListModelVersionsRequest, + model_service.ListModelVersionCheckpointsRequest, dict, ], ) -def test_list_model_versions_rest_call_success(request_type): +def test_list_model_version_checkpoints_rest_call_success(request_type): client = ModelServiceClient( credentials=ga_credentials.AnonymousCredentials(), transport="rest" ) @@ -13597,7 +14685,7 @@ def test_list_model_versions_rest_call_success(request_type): # Mock the http request call within the method and fake a response. with mock.patch.object(type(client.transport._session), "request") as req: # Designate an appropriate value for the returned response. - return_value = model_service.ListModelVersionsResponse( + return_value = model_service.ListModelVersionCheckpointsResponse( next_page_token="next_page_token_value", ) @@ -13606,20 +14694,22 @@ def test_list_model_versions_rest_call_success(request_type): response_value.status_code = 200 # Convert return value to protobuf type - return_value = model_service.ListModelVersionsResponse.pb(return_value) + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) json_return_value = json_format.MessageToJson(return_value) response_value.content = json_return_value.encode("UTF-8") req.return_value = response_value req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - response = client.list_model_versions(request) + response = client.list_model_version_checkpoints(request) # Establish that the response is the type that we expect. - assert isinstance(response, pagers.ListModelVersionsPager) + assert isinstance(response, pagers.ListModelVersionCheckpointsPager) assert response.next_page_token == "next_page_token_value" @pytest.mark.parametrize("null_interceptor", [True, False]) -def test_list_model_versions_rest_interceptors(null_interceptor): +def test_list_model_version_checkpoints_rest_interceptors(null_interceptor): transport = transports.ModelServiceRestTransport( credentials=ga_credentials.AnonymousCredentials(), interceptor=None @@ -13633,14 +14723,18 @@ def test_list_model_versions_rest_interceptors(null_interceptor): ) as req, mock.patch.object( path_template, "transcode" ) as transcode, mock.patch.object( - transports.ModelServiceRestInterceptor, "post_list_model_versions" + transports.ModelServiceRestInterceptor, "post_list_model_version_checkpoints" ) as post, mock.patch.object( - transports.ModelServiceRestInterceptor, "pre_list_model_versions" + transports.ModelServiceRestInterceptor, + "post_list_model_version_checkpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.ModelServiceRestInterceptor, "pre_list_model_version_checkpoints" ) as pre: pre.assert_not_called() post.assert_not_called() - pb_message = model_service.ListModelVersionsRequest.pb( - model_service.ListModelVersionsRequest() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionCheckpointsRequest.pb( + model_service.ListModelVersionCheckpointsRequest() ) transcode.return_value = { "method": "post", @@ -13652,20 +14746,24 @@ def test_list_model_versions_rest_interceptors(null_interceptor): req.return_value = mock.Mock() req.return_value.status_code = 200 req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} - return_value = model_service.ListModelVersionsResponse.to_json( - model_service.ListModelVersionsResponse() + return_value = model_service.ListModelVersionCheckpointsResponse.to_json( + model_service.ListModelVersionCheckpointsResponse() ) req.return_value.content = return_value - request = model_service.ListModelVersionsRequest() + request = model_service.ListModelVersionCheckpointsRequest() metadata = [ ("key", "val"), ("cephalopod", "squid"), ] pre.return_value = request, metadata - post.return_value = model_service.ListModelVersionsResponse() + post.return_value = model_service.ListModelVersionCheckpointsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionCheckpointsResponse(), + metadata, + ) - client.list_model_versions( + client.list_model_version_checkpoints( request, metadata=[ ("key", "val"), @@ -13675,6 +14773,7 @@ def test_list_model_versions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_model_rest_bad_request(request_type=model_service.UpdateModelRequest): @@ -13727,6 +14826,7 @@ def test_update_model_rest_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -13756,10 +14856,23 @@ def test_update_model_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [{"name": "name_value", "value": "value_value"}], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -13910,6 +15023,7 @@ def get_message_fields(field): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -13944,6 +15058,7 @@ def get_message_fields(field): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -13979,10 +15094,13 @@ def test_update_model_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_update_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_update_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_update_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateModelRequest.pb( model_service.UpdateModelRequest() ) @@ -14006,6 +15124,7 @@ def test_update_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model.Model() + post_with_metadata.return_value = gca_model.Model(), metadata client.update_model( request, @@ -14017,6 +15136,7 @@ def test_update_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_explanation_dataset_rest_bad_request( @@ -14097,10 +15217,14 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_update_explanation_dataset" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_update_explanation_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_update_explanation_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateExplanationDatasetRequest.pb( model_service.UpdateExplanationDatasetRequest() ) @@ -14124,6 +15248,7 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_explanation_dataset( request, @@ -14135,6 +15260,7 @@ def test_update_explanation_dataset_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_rest_bad_request(request_type=model_service.DeleteModelRequest): @@ -14213,10 +15339,13 @@ def test_delete_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_delete_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_delete_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_delete_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelRequest.pb( model_service.DeleteModelRequest() ) @@ -14240,6 +15369,7 @@ def test_delete_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model( request, @@ -14251,6 +15381,7 @@ def test_delete_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_model_version_rest_bad_request( @@ -14331,10 +15462,14 @@ def test_delete_model_version_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_delete_model_version" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_delete_model_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_delete_model_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelVersionRequest.pb( model_service.DeleteModelVersionRequest() ) @@ -14358,6 +15493,7 @@ def test_delete_model_version_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_model_version( request, @@ -14369,6 +15505,7 @@ def test_delete_model_version_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_merge_version_aliases_rest_bad_request( @@ -14422,6 +15559,7 @@ def test_merge_version_aliases_rest_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -14456,6 +15594,7 @@ def test_merge_version_aliases_rest_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -14491,10 +15630,14 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_merge_version_aliases" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_merge_version_aliases_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_merge_version_aliases" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.MergeVersionAliasesRequest.pb( model_service.MergeVersionAliasesRequest() ) @@ -14518,6 +15661,7 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata client.merge_version_aliases( request, @@ -14529,6 +15673,7 @@ def test_merge_version_aliases_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_model_rest_bad_request(request_type=model_service.ExportModelRequest): @@ -14607,10 +15752,13 @@ def test_export_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_export_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_export_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_export_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ExportModelRequest.pb( model_service.ExportModelRequest() ) @@ -14634,6 +15782,7 @@ def test_export_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.export_model( request, @@ -14645,6 +15794,7 @@ def test_export_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_copy_model_rest_bad_request(request_type=model_service.CopyModelRequest): @@ -14723,10 +15873,13 @@ def test_copy_model_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ModelServiceRestInterceptor, "post_copy_model" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, "post_copy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_copy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.CopyModelRequest.pb(model_service.CopyModelRequest()) transcode.return_value = { "method": "post", @@ -14748,6 +15901,7 @@ def test_copy_model_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.copy_model( request, @@ -14759,6 +15913,7 @@ def test_copy_model_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_model_evaluation_rest_bad_request( @@ -14849,10 +16004,14 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_import_model_evaluation" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_import_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_import_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ImportModelEvaluationRequest.pb( model_service.ImportModelEvaluationRequest() ) @@ -14878,6 +16037,10 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model_evaluation.ModelEvaluation() + post_with_metadata.return_value = ( + gca_model_evaluation.ModelEvaluation(), + metadata, + ) client.import_model_evaluation( request, @@ -14889,6 +16052,7 @@ def test_import_model_evaluation_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_import_model_evaluation_slices_rest_bad_request( @@ -14982,11 +16146,15 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor transports.ModelServiceRestInterceptor, "post_batch_import_model_evaluation_slices", ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_batch_import_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_batch_import_model_evaluation_slices", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportModelEvaluationSlicesRequest.pb( model_service.BatchImportModelEvaluationSlicesRequest() ) @@ -15012,6 +16180,10 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = model_service.BatchImportModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.BatchImportModelEvaluationSlicesResponse(), + metadata, + ) client.batch_import_model_evaluation_slices( request, @@ -15023,6 +16195,7 @@ def test_batch_import_model_evaluation_slices_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_import_evaluated_annotations_rest_bad_request( @@ -15114,10 +16287,14 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): transports.ModelServiceRestInterceptor, "post_batch_import_evaluated_annotations", ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_batch_import_evaluated_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_batch_import_evaluated_annotations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportEvaluatedAnnotationsRequest.pb( model_service.BatchImportEvaluatedAnnotationsRequest() ) @@ -15143,6 +16320,10 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.BatchImportEvaluatedAnnotationsResponse() + post_with_metadata.return_value = ( + model_service.BatchImportEvaluatedAnnotationsResponse(), + metadata, + ) client.batch_import_evaluated_annotations( request, @@ -15154,6 +16335,7 @@ def test_batch_import_evaluated_annotations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_evaluation_rest_bad_request( @@ -15248,10 +16430,14 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model_evaluation" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_get_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationRequest.pb( model_service.GetModelEvaluationRequest() ) @@ -15277,6 +16463,7 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation.ModelEvaluation() + post_with_metadata.return_value = model_evaluation.ModelEvaluation(), metadata client.get_model_evaluation( request, @@ -15288,6 +16475,7 @@ def test_get_model_evaluation_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_evaluations_rest_bad_request( @@ -15372,10 +16560,14 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_list_model_evaluations" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_list_model_evaluations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_list_model_evaluations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationsRequest.pb( model_service.ListModelEvaluationsRequest() ) @@ -15401,6 +16593,10 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationsResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationsResponse(), + metadata, + ) client.list_model_evaluations( request, @@ -15412,6 +16608,7 @@ def test_list_model_evaluations_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_model_evaluation_slice_rest_bad_request( @@ -15502,10 +16699,14 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_get_model_evaluation_slice" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_get_model_evaluation_slice_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_get_model_evaluation_slice" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationSliceRequest.pb( model_service.GetModelEvaluationSliceRequest() ) @@ -15531,6 +16732,10 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation_slice.ModelEvaluationSlice() + post_with_metadata.return_value = ( + model_evaluation_slice.ModelEvaluationSlice(), + metadata, + ) client.get_model_evaluation_slice( request, @@ -15542,6 +16747,7 @@ def test_get_model_evaluation_slice_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_model_evaluation_slices_rest_bad_request( @@ -15630,10 +16836,14 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ModelServiceRestInterceptor, "post_list_model_evaluation_slices" ) as post, mock.patch.object( + transports.ModelServiceRestInterceptor, + "post_list_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ModelServiceRestInterceptor, "pre_list_model_evaluation_slices" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationSlicesRequest.pb( model_service.ListModelEvaluationSlicesRequest() ) @@ -15659,6 +16869,10 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationSlicesResponse(), + metadata, + ) client.list_model_evaluation_slices( request, @@ -15670,6 +16884,7 @@ def test_list_model_evaluation_slices_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -16386,6 +17601,28 @@ def test_list_model_versions_empty_call_rest(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +def test_list_model_version_checkpoints_empty_call_rest(): + client = ModelServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. def test_update_model_empty_call_rest(): @@ -16808,10 +18045,13 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_upload_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_upload_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_upload_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UploadModelRequest.pb( model_service.UploadModelRequest() ) @@ -16835,6 +18075,7 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.upload_model( request, @@ -16846,6 +18087,7 @@ async def test_upload_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16908,6 +18150,7 @@ async def test_get_model_rest_asyncio_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -16944,6 +18187,7 @@ async def test_get_model_rest_asyncio_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -16984,10 +18228,13 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_get_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelRequest.pb(model_service.GetModelRequest()) transcode.return_value = { "method": "post", @@ -17009,6 +18256,7 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata await client.get_model( request, @@ -17020,6 +18268,7 @@ async def test_get_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17120,10 +18369,13 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_models" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_list_models_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_models" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelsRequest.pb( model_service.ListModelsRequest() ) @@ -17149,6 +18401,7 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelsResponse() + post_with_metadata.return_value = model_service.ListModelsResponse(), metadata await client.list_models( request, @@ -17160,6 +18413,7 @@ async def test_list_models_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17260,10 +18514,14 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_versions" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_versions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_versions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelVersionsRequest.pb( model_service.ListModelVersionsRequest() ) @@ -17289,6 +18547,10 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_service.ListModelVersionsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionsResponse(), + metadata, + ) await client.list_model_versions( request, @@ -17300,6 +18562,162 @@ async def test_list_model_versions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_rest_asyncio_bad_request( + request_type=model_service.ListModelVersionCheckpointsRequest, +): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(AsyncAuthorizedSession, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.read = mock.AsyncMock(return_value=b"{}") + response_value.status_code = 400 + response_value.request = mock.Mock() + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + await client.list_model_version_checkpoints(request) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_type", + [ + model_service.ListModelVersionCheckpointsRequest, + dict, + ], +) +async def test_list_model_version_checkpoints_rest_asyncio_call_success(request_type): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), transport="rest_asyncio" + ) + + # send a request that will satisfy transcoding + request_init = {"name": "projects/sample1/locations/sample2/models/sample3"} + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = model_service.ListModelVersionCheckpointsResponse( + next_page_token="next_page_token_value", + ) + + # Wrap the value into a proper Response obj + response_value = mock.Mock() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = model_service.ListModelVersionCheckpointsResponse.pb( + return_value + ) + json_return_value = json_format.MessageToJson(return_value) + response_value.read = mock.AsyncMock( + return_value=json_return_value.encode("UTF-8") + ) + req.return_value = response_value + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + response = await client.list_model_version_checkpoints(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, pagers.ListModelVersionCheckpointsAsyncPager) + assert response.next_page_token == "next_page_token_value" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("null_interceptor", [True, False]) +async def test_list_model_version_checkpoints_rest_asyncio_interceptors( + null_interceptor, +): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + transport = transports.AsyncModelServiceRestTransport( + credentials=async_anonymous_credentials(), + interceptor=None + if null_interceptor + else transports.AsyncModelServiceRestInterceptor(), + ) + client = ModelServiceAsyncClient(transport=transport) + + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_version_checkpoints", + ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_version_checkpoints_with_metadata", + ) as post_with_metadata, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "pre_list_model_version_checkpoints", + ) as pre: + pre.assert_not_called() + post.assert_not_called() + post_with_metadata.assert_not_called() + pb_message = model_service.ListModelVersionCheckpointsRequest.pb( + model_service.ListModelVersionCheckpointsRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = mock.Mock() + req.return_value.status_code = 200 + req.return_value.headers = {"header-1": "value-1", "header-2": "value-2"} + return_value = model_service.ListModelVersionCheckpointsResponse.to_json( + model_service.ListModelVersionCheckpointsResponse() + ) + req.return_value.read = mock.AsyncMock(return_value=return_value) + + request = model_service.ListModelVersionCheckpointsRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = model_service.ListModelVersionCheckpointsResponse() + post_with_metadata.return_value = ( + model_service.ListModelVersionCheckpointsResponse(), + metadata, + ) + + await client.list_model_version_checkpoints( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17363,6 +18781,7 @@ async def test_update_model_rest_asyncio_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -17392,10 +18811,23 @@ async def test_update_model_rest_asyncio_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [{"name": "name_value", "value": "value_value"}], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -17546,6 +18978,7 @@ def get_message_fields(field): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -17582,6 +19015,7 @@ def get_message_fields(field): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -17622,10 +19056,13 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_update_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_update_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_update_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateModelRequest.pb( model_service.UpdateModelRequest() ) @@ -17649,6 +19086,7 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_model.Model() + post_with_metadata.return_value = gca_model.Model(), metadata await client.update_model( request, @@ -17660,6 +19098,7 @@ async def test_update_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17756,10 +19195,14 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_update_explanation_dataset" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_update_explanation_dataset_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_update_explanation_dataset" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.UpdateExplanationDatasetRequest.pb( model_service.UpdateExplanationDatasetRequest() ) @@ -17783,6 +19226,7 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_explanation_dataset( request, @@ -17794,6 +19238,7 @@ async def test_update_explanation_dataset_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17890,10 +19335,13 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_delete_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_delete_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_delete_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelRequest.pb( model_service.DeleteModelRequest() ) @@ -17917,6 +19365,7 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model( request, @@ -17928,6 +19377,7 @@ async def test_delete_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18024,10 +19474,14 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_delete_model_version" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_delete_model_version_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_delete_model_version" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.DeleteModelVersionRequest.pb( model_service.DeleteModelVersionRequest() ) @@ -18051,6 +19505,7 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_model_version( request, @@ -18062,6 +19517,7 @@ async def test_delete_model_version_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18124,6 +19580,7 @@ async def test_merge_version_aliases_rest_asyncio_call_success(request_type): display_name="display_name_value", description="description_value", version_description="version_description_value", + default_checkpoint_id="default_checkpoint_id_value", metadata_schema_uri="metadata_schema_uri_value", training_pipeline="training_pipeline_value", artifact_uri="artifact_uri_value", @@ -18160,6 +19617,7 @@ async def test_merge_version_aliases_rest_asyncio_call_success(request_type): assert response.display_name == "display_name_value" assert response.description == "description_value" assert response.version_description == "version_description_value" + assert response.default_checkpoint_id == "default_checkpoint_id_value" assert response.metadata_schema_uri == "metadata_schema_uri_value" assert response.training_pipeline == "training_pipeline_value" assert response.artifact_uri == "artifact_uri_value" @@ -18200,10 +19658,14 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_merge_version_aliases" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_merge_version_aliases_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_merge_version_aliases" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.MergeVersionAliasesRequest.pb( model_service.MergeVersionAliasesRequest() ) @@ -18227,6 +19689,7 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = model.Model() + post_with_metadata.return_value = model.Model(), metadata await client.merge_version_aliases( request, @@ -18238,6 +19701,7 @@ async def test_merge_version_aliases_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18334,10 +19798,13 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_export_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_export_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_export_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ExportModelRequest.pb( model_service.ExportModelRequest() ) @@ -18361,6 +19828,7 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.export_model( request, @@ -18372,6 +19840,7 @@ async def test_export_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18468,10 +19937,13 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_copy_model" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, "post_copy_model_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_copy_model" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.CopyModelRequest.pb(model_service.CopyModelRequest()) transcode.return_value = { "method": "post", @@ -18493,6 +19965,7 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.copy_model( request, @@ -18504,6 +19977,7 @@ async def test_copy_model_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18610,10 +20084,14 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_import_model_evaluation" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_import_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_import_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ImportModelEvaluationRequest.pb( model_service.ImportModelEvaluationRequest() ) @@ -18639,6 +20117,10 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = gca_model_evaluation.ModelEvaluation() + post_with_metadata.return_value = ( + gca_model_evaluation.ModelEvaluation(), + metadata, + ) await client.import_model_evaluation( request, @@ -18650,6 +20132,7 @@ async def test_import_model_evaluation_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18763,11 +20246,15 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( transports.AsyncModelServiceRestInterceptor, "post_batch_import_model_evaluation_slices", ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_batch_import_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_batch_import_model_evaluation_slices", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportModelEvaluationSlicesRequest.pb( model_service.BatchImportModelEvaluationSlicesRequest() ) @@ -18793,6 +20280,10 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = model_service.BatchImportModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.BatchImportModelEvaluationSlicesResponse(), + metadata, + ) await client.batch_import_model_evaluation_slices( request, @@ -18804,6 +20295,7 @@ async def test_batch_import_model_evaluation_slices_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -18915,11 +20407,15 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( transports.AsyncModelServiceRestInterceptor, "post_batch_import_evaluated_annotations", ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_batch_import_evaluated_annotations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_batch_import_evaluated_annotations", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.BatchImportEvaluatedAnnotationsRequest.pb( model_service.BatchImportEvaluatedAnnotationsRequest() ) @@ -18945,6 +20441,10 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = model_service.BatchImportEvaluatedAnnotationsResponse() + post_with_metadata.return_value = ( + model_service.BatchImportEvaluatedAnnotationsResponse(), + metadata, + ) await client.batch_import_evaluated_annotations( request, @@ -18956,6 +20456,7 @@ async def test_batch_import_evaluated_annotations_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19066,10 +20567,14 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model_evaluation" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_get_model_evaluation_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model_evaluation" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationRequest.pb( model_service.GetModelEvaluationRequest() ) @@ -19095,6 +20600,7 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = model_evaluation.ModelEvaluation() + post_with_metadata.return_value = model_evaluation.ModelEvaluation(), metadata await client.get_model_evaluation( request, @@ -19106,6 +20612,7 @@ async def test_get_model_evaluation_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19206,10 +20713,14 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_evaluations" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_evaluations_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_evaluations" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationsRequest.pb( model_service.ListModelEvaluationsRequest() ) @@ -19235,6 +20746,10 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationsResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationsResponse(), + metadata, + ) await client.list_model_evaluations( request, @@ -19246,6 +20761,7 @@ async def test_list_model_evaluations_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19352,10 +20868,14 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_get_model_evaluation_slice" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_get_model_evaluation_slice_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_get_model_evaluation_slice" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.GetModelEvaluationSliceRequest.pb( model_service.GetModelEvaluationSliceRequest() ) @@ -19381,6 +20901,10 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = model_evaluation_slice.ModelEvaluationSlice() + post_with_metadata.return_value = ( + model_evaluation_slice.ModelEvaluationSlice(), + metadata, + ) await client.get_model_evaluation_slice( request, @@ -19392,6 +20916,7 @@ async def test_get_model_evaluation_slice_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -19496,10 +21021,14 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter ) as transcode, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "post_list_model_evaluation_slices" ) as post, mock.patch.object( + transports.AsyncModelServiceRestInterceptor, + "post_list_model_evaluation_slices_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncModelServiceRestInterceptor, "pre_list_model_evaluation_slices" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = model_service.ListModelEvaluationSlicesRequest.pb( model_service.ListModelEvaluationSlicesRequest() ) @@ -19525,6 +21054,10 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = model_service.ListModelEvaluationSlicesResponse() + post_with_metadata.return_value = ( + model_service.ListModelEvaluationSlicesResponse(), + metadata, + ) await client.list_model_evaluation_slices( request, @@ -19536,6 +21069,7 @@ async def test_list_model_evaluation_slices_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -20388,6 +21922,33 @@ async def test_list_model_versions_empty_call_rest_asyncio(): assert args[0] == request_msg +# This test is a coverage failsafe to make sure that totally empty calls, +# i.e. request == None and no flattened fields passed, work. +@pytest.mark.asyncio +async def test_list_model_version_checkpoints_empty_call_rest_asyncio(): + if not HAS_ASYNC_REST_EXTRA: + pytest.skip( + "the library must be installed with the `async_rest` extra to test this feature." + ) + client = ModelServiceAsyncClient( + credentials=async_anonymous_credentials(), + transport="rest_asyncio", + ) + + # Mock the actual call, and fake the request. + with mock.patch.object( + type(client.transport.list_model_version_checkpoints), "__call__" + ) as call: + await client.list_model_version_checkpoints(request=None) + + # Establish that the underlying stub method was called. + call.assert_called() + _, args, _ = call.mock_calls[0] + request_msg = model_service.ListModelVersionCheckpointsRequest() + + assert args[0] == request_msg + + # This test is a coverage failsafe to make sure that totally empty calls, # i.e. request == None and no flattened fields passed, work. @pytest.mark.asyncio @@ -20830,6 +22391,7 @@ def test_model_service_base_transport(): "get_model", "list_models", "list_model_versions", + "list_model_version_checkpoints", "update_model", "update_explanation_dataset", "delete_model", @@ -21126,6 +22688,9 @@ def test_model_service_client_transport_session_collision(transport_name): session1 = client1.transport.list_model_versions._session session2 = client2.transport.list_model_versions._session assert session1 != session2 + session1 = client1.transport.list_model_version_checkpoints._session + session2 = client2.transport.list_model_version_checkpoints._session + assert session1 != session2 session1 = client1.transport.update_model._session session2 = client2.transport.update_model._session assert session1 != session2 diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_notebook_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_notebook_service.py index 39cf8eef7c..124db731cf 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_notebook_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_notebook_service.py @@ -76,6 +76,7 @@ from google.cloud.aiplatform_v1beta1.services.notebook_service import transports from google.cloud.aiplatform_v1beta1.types import accelerator_type from google.cloud.aiplatform_v1beta1.types import encryption_spec +from google.cloud.aiplatform_v1beta1.types import env_var from google.cloud.aiplatform_v1beta1.types import job_state from google.cloud.aiplatform_v1beta1.types import machine_resources from google.cloud.aiplatform_v1beta1.types import network_spec @@ -91,6 +92,7 @@ ) from google.cloud.aiplatform_v1beta1.types import notebook_runtime_template_ref from google.cloud.aiplatform_v1beta1.types import notebook_service +from google.cloud.aiplatform_v1beta1.types import notebook_software_config from google.cloud.aiplatform_v1beta1.types import operation as gca_operation from google.cloud.aiplatform_v1beta1.types import reservation_affinity from google.cloud.location import locations_pb2 @@ -108,6 +110,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -365,6 +375,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = NotebookServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = NotebookServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -11920,6 +11973,14 @@ def test_create_notebook_runtime_template_rest_call_success(request_type): "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -12034,11 +12095,15 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_create_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_create_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_create_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookRuntimeTemplateRequest.pb( notebook_service.CreateNotebookRuntimeTemplateRequest() ) @@ -12062,6 +12127,7 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_notebook_runtime_template( request, @@ -12073,6 +12139,7 @@ def test_create_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_runtime_template_rest_bad_request( @@ -12178,10 +12245,14 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_runtime_template" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_runtime_template" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeTemplateRequest.pb( notebook_service.GetNotebookRuntimeTemplateRequest() ) @@ -12207,6 +12278,10 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) client.get_notebook_runtime_template( request, @@ -12218,6 +12293,7 @@ def test_get_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_runtime_templates_rest_bad_request( @@ -12305,10 +12381,14 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_list_notebook_runtime_templates", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_runtime_templates_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_runtime_templates" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimeTemplatesRequest.pb( notebook_service.ListNotebookRuntimeTemplatesRequest() ) @@ -12334,6 +12414,10 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimeTemplatesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimeTemplatesResponse(), + metadata, + ) client.list_notebook_runtime_templates( request, @@ -12345,6 +12429,7 @@ def test_list_notebook_runtime_templates_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_runtime_template_rest_bad_request( @@ -12430,11 +12515,15 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_delete_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeTemplateRequest.pb( notebook_service.DeleteNotebookRuntimeTemplateRequest() ) @@ -12458,6 +12547,7 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_runtime_template( request, @@ -12469,6 +12559,7 @@ def test_delete_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_notebook_runtime_template_rest_bad_request( @@ -12557,6 +12648,14 @@ def test_update_notebook_runtime_template_rest_call_success(request_type): "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -12692,11 +12791,15 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): transports.NotebookServiceRestInterceptor, "post_update_notebook_runtime_template", ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_update_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_update_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpdateNotebookRuntimeTemplateRequest.pb( notebook_service.UpdateNotebookRuntimeTemplateRequest() ) @@ -12722,6 +12825,10 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) client.update_notebook_runtime_template( request, @@ -12733,6 +12840,7 @@ def test_update_notebook_runtime_template_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_assign_notebook_runtime_rest_bad_request( @@ -12813,10 +12921,14 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_assign_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_assign_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_assign_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.AssignNotebookRuntimeRequest.pb( notebook_service.AssignNotebookRuntimeRequest() ) @@ -12840,6 +12952,7 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.assign_notebook_runtime( request, @@ -12851,6 +12964,7 @@ def test_assign_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_runtime_rest_bad_request( @@ -12970,10 +13084,14 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeRequest.pb( notebook_service.GetNotebookRuntimeRequest() ) @@ -12999,6 +13117,7 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntime() + post_with_metadata.return_value = notebook_runtime.NotebookRuntime(), metadata client.get_notebook_runtime( request, @@ -13010,6 +13129,7 @@ def test_get_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_runtimes_rest_bad_request( @@ -13094,10 +13214,14 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_list_notebook_runtimes" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_runtimes_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_runtimes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimesRequest.pb( notebook_service.ListNotebookRuntimesRequest() ) @@ -13123,6 +13247,10 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimesResponse(), + metadata, + ) client.list_notebook_runtimes( request, @@ -13134,6 +13262,7 @@ def test_list_notebook_runtimes_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_runtime_rest_bad_request( @@ -13218,10 +13347,14 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_delete_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeRequest.pb( notebook_service.DeleteNotebookRuntimeRequest() ) @@ -13245,6 +13378,7 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_runtime( request, @@ -13256,6 +13390,7 @@ def test_delete_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upgrade_notebook_runtime_rest_bad_request( @@ -13340,10 +13475,14 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_upgrade_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_upgrade_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_upgrade_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpgradeNotebookRuntimeRequest.pb( notebook_service.UpgradeNotebookRuntimeRequest() ) @@ -13367,6 +13506,7 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.upgrade_notebook_runtime( request, @@ -13378,6 +13518,7 @@ def test_upgrade_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_start_notebook_runtime_rest_bad_request( @@ -13462,10 +13603,14 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_start_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_start_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_start_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StartNotebookRuntimeRequest.pb( notebook_service.StartNotebookRuntimeRequest() ) @@ -13489,6 +13634,7 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.start_notebook_runtime( request, @@ -13500,6 +13646,7 @@ def test_start_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stop_notebook_runtime_rest_bad_request( @@ -13584,10 +13731,14 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_stop_notebook_runtime" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_stop_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_stop_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StopNotebookRuntimeRequest.pb( notebook_service.StopNotebookRuntimeRequest() ) @@ -13611,6 +13762,7 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.stop_notebook_runtime( request, @@ -13622,6 +13774,7 @@ def test_stop_notebook_runtime_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_notebook_execution_job_rest_bad_request( @@ -13828,10 +13981,14 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_create_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_create_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_create_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookExecutionJobRequest.pb( notebook_service.CreateNotebookExecutionJobRequest() ) @@ -13855,6 +14012,7 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_notebook_execution_job( request, @@ -13866,6 +14024,7 @@ def test_create_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_notebook_execution_job_rest_bad_request( @@ -13965,10 +14124,14 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_get_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_get_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_get_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookExecutionJobRequest.pb( notebook_service.GetNotebookExecutionJobRequest() ) @@ -13994,6 +14157,10 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_execution_job.NotebookExecutionJob() + post_with_metadata.return_value = ( + notebook_execution_job.NotebookExecutionJob(), + metadata, + ) client.get_notebook_execution_job( request, @@ -14005,6 +14172,7 @@ def test_get_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_notebook_execution_jobs_rest_bad_request( @@ -14091,10 +14259,14 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.NotebookServiceRestInterceptor, "post_list_notebook_execution_jobs" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_list_notebook_execution_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_list_notebook_execution_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookExecutionJobsRequest.pb( notebook_service.ListNotebookExecutionJobsRequest() ) @@ -14120,6 +14292,10 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookExecutionJobsResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookExecutionJobsResponse(), + metadata, + ) client.list_notebook_execution_jobs( request, @@ -14131,6 +14307,7 @@ def test_list_notebook_execution_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_notebook_execution_job_rest_bad_request( @@ -14215,10 +14392,14 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.NotebookServiceRestInterceptor, "post_delete_notebook_execution_job" ) as post, mock.patch.object( + transports.NotebookServiceRestInterceptor, + "post_delete_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.NotebookServiceRestInterceptor, "pre_delete_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookExecutionJobRequest.pb( notebook_service.DeleteNotebookExecutionJobRequest() ) @@ -14242,6 +14423,7 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_notebook_execution_job( request, @@ -14253,6 +14435,7 @@ def test_delete_notebook_execution_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -15354,6 +15537,14 @@ async def test_create_notebook_runtime_template_rest_asyncio_call_success(reques "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -15477,11 +15668,15 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_create_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_create_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_create_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookRuntimeTemplateRequest.pb( notebook_service.CreateNotebookRuntimeTemplateRequest() ) @@ -15505,6 +15700,7 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_notebook_runtime_template( request, @@ -15516,6 +15712,7 @@ async def test_create_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15640,11 +15837,15 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeTemplateRequest.pb( notebook_service.GetNotebookRuntimeTemplateRequest() ) @@ -15670,6 +15871,10 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) await client.get_notebook_runtime_template( request, @@ -15681,6 +15886,7 @@ async def test_get_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15786,11 +15992,15 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_runtime_templates", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_runtime_templates_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_runtime_templates", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimeTemplatesRequest.pb( notebook_service.ListNotebookRuntimeTemplatesRequest() ) @@ -15816,6 +16026,10 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimeTemplatesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimeTemplatesResponse(), + metadata, + ) await client.list_notebook_runtime_templates( request, @@ -15827,6 +16041,7 @@ async def test_list_notebook_runtime_templates_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -15930,11 +16145,15 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeTemplateRequest.pb( notebook_service.DeleteNotebookRuntimeTemplateRequest() ) @@ -15958,6 +16177,7 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_runtime_template( request, @@ -15969,6 +16189,7 @@ async def test_delete_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16066,6 +16287,14 @@ async def test_update_notebook_runtime_template_rest_asyncio_call_success(reques "shielded_vm_config": {"enable_secure_boot": True}, "network_tags": ["network_tags_value1", "network_tags_value2"], "encryption_spec": {"kms_key_name": "kms_key_name_value"}, + "software_config": { + "env": [{"name": "name_value", "value": "value_value"}], + "post_startup_script_config": { + "post_startup_script": "post_startup_script_value", + "post_startup_script_url": "post_startup_script_url_value", + "post_startup_script_behavior": 1, + }, + }, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -16210,11 +16439,15 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_update_notebook_runtime_template", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_update_notebook_runtime_template_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_update_notebook_runtime_template", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpdateNotebookRuntimeTemplateRequest.pb( notebook_service.UpdateNotebookRuntimeTemplateRequest() ) @@ -16240,6 +16473,10 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntimeTemplate() + post_with_metadata.return_value = ( + notebook_runtime.NotebookRuntimeTemplate(), + metadata, + ) await client.update_notebook_runtime_template( request, @@ -16251,6 +16488,7 @@ async def test_update_notebook_runtime_template_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16347,10 +16585,14 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_assign_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_assign_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_assign_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.AssignNotebookRuntimeRequest.pb( notebook_service.AssignNotebookRuntimeRequest() ) @@ -16374,6 +16616,7 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.assign_notebook_runtime( request, @@ -16385,6 +16628,7 @@ async def test_assign_notebook_runtime_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16520,10 +16764,14 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookRuntimeRequest.pb( notebook_service.GetNotebookRuntimeRequest() ) @@ -16549,6 +16797,7 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = notebook_runtime.NotebookRuntime() + post_with_metadata.return_value = notebook_runtime.NotebookRuntime(), metadata await client.get_notebook_runtime( request, @@ -16560,6 +16809,7 @@ async def test_get_notebook_runtime_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16660,10 +16910,14 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_runtimes" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_runtimes_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_runtimes" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookRuntimesRequest.pb( notebook_service.ListNotebookRuntimesRequest() ) @@ -16689,6 +16943,10 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookRuntimesResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookRuntimesResponse(), + metadata, + ) await client.list_notebook_runtimes( request, @@ -16700,6 +16958,7 @@ async def test_list_notebook_runtimes_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16800,10 +17059,14 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookRuntimeRequest.pb( notebook_service.DeleteNotebookRuntimeRequest() ) @@ -16827,6 +17090,7 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_runtime( request, @@ -16838,6 +17102,7 @@ async def test_delete_notebook_runtime_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -16938,10 +17203,14 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_upgrade_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_upgrade_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_upgrade_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.UpgradeNotebookRuntimeRequest.pb( notebook_service.UpgradeNotebookRuntimeRequest() ) @@ -16965,6 +17234,7 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.upgrade_notebook_runtime( request, @@ -16976,6 +17246,7 @@ async def test_upgrade_notebook_runtime_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17076,10 +17347,14 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_start_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_start_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_start_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StartNotebookRuntimeRequest.pb( notebook_service.StartNotebookRuntimeRequest() ) @@ -17103,6 +17378,7 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.start_notebook_runtime( request, @@ -17114,6 +17390,7 @@ async def test_start_notebook_runtime_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17214,10 +17491,14 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) ), mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "post_stop_notebook_runtime" ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_stop_notebook_runtime_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_stop_notebook_runtime" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.StopNotebookRuntimeRequest.pb( notebook_service.StopNotebookRuntimeRequest() ) @@ -17241,6 +17522,7 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.stop_notebook_runtime( request, @@ -17252,6 +17534,7 @@ async def test_stop_notebook_runtime_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17477,11 +17760,15 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_create_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_create_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_create_notebook_execution_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.CreateNotebookExecutionJobRequest.pb( notebook_service.CreateNotebookExecutionJobRequest() ) @@ -17505,6 +17792,7 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_notebook_execution_job( request, @@ -17516,6 +17804,7 @@ async def test_create_notebook_execution_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17632,10 +17921,14 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce transports.AsyncNotebookServiceRestInterceptor, "post_get_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_get_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_get_notebook_execution_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.GetNotebookExecutionJobRequest.pb( notebook_service.GetNotebookExecutionJobRequest() ) @@ -17661,6 +17954,10 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = notebook_execution_job.NotebookExecutionJob() + post_with_metadata.return_value = ( + notebook_execution_job.NotebookExecutionJob(), + metadata, + ) await client.get_notebook_execution_job( request, @@ -17672,6 +17969,7 @@ async def test_get_notebook_execution_job_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17775,11 +18073,15 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter transports.AsyncNotebookServiceRestInterceptor, "post_list_notebook_execution_jobs", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_list_notebook_execution_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_list_notebook_execution_jobs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.ListNotebookExecutionJobsRequest.pb( notebook_service.ListNotebookExecutionJobsRequest() ) @@ -17805,6 +18107,10 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = notebook_service.ListNotebookExecutionJobsResponse() + post_with_metadata.return_value = ( + notebook_service.ListNotebookExecutionJobsResponse(), + metadata, + ) await client.list_notebook_execution_jobs( request, @@ -17816,6 +18122,7 @@ async def test_list_notebook_execution_jobs_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -17919,11 +18226,15 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( transports.AsyncNotebookServiceRestInterceptor, "post_delete_notebook_execution_job", ) as post, mock.patch.object( + transports.AsyncNotebookServiceRestInterceptor, + "post_delete_notebook_execution_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncNotebookServiceRestInterceptor, "pre_delete_notebook_execution_job", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = notebook_service.DeleteNotebookExecutionJobRequest.pb( notebook_service.DeleteNotebookExecutionJobRequest() ) @@ -17947,6 +18258,7 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_notebook_execution_job( request, @@ -17958,6 +18270,7 @@ async def test_delete_notebook_execution_job_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_persistent_resource_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_persistent_resource_service.py index 745f1cd69b..51bdc89aa0 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_persistent_resource_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_persistent_resource_service.py @@ -101,6 +101,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -379,6 +387,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PersistentResourceServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PersistentResourceServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -5441,11 +5492,15 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_create_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_create_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_create_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.CreatePersistentResourceRequest.pb( persistent_resource_service.CreatePersistentResourceRequest() ) @@ -5469,6 +5524,7 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_persistent_resource( request, @@ -5480,6 +5536,7 @@ def test_create_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_persistent_resource_rest_bad_request( @@ -5581,11 +5638,15 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_get_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_get_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_get_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.GetPersistentResourceRequest.pb( persistent_resource_service.GetPersistentResourceRequest() ) @@ -5611,6 +5672,10 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = persistent_resource.PersistentResource() + post_with_metadata.return_value = ( + persistent_resource.PersistentResource(), + metadata, + ) client.get_persistent_resource( request, @@ -5622,6 +5687,7 @@ def test_get_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_persistent_resources_rest_bad_request( @@ -5709,11 +5775,15 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_list_persistent_resources", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_list_persistent_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_list_persistent_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.ListPersistentResourcesRequest.pb( persistent_resource_service.ListPersistentResourcesRequest() ) @@ -5743,6 +5813,10 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): post.return_value = ( persistent_resource_service.ListPersistentResourcesResponse() ) + post_with_metadata.return_value = ( + persistent_resource_service.ListPersistentResourcesResponse(), + metadata, + ) client.list_persistent_resources( request, @@ -5754,6 +5828,7 @@ def test_list_persistent_resources_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_persistent_resource_rest_bad_request( @@ -5839,11 +5914,15 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_delete_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_delete_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_delete_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.DeletePersistentResourceRequest.pb( persistent_resource_service.DeletePersistentResourceRequest() ) @@ -5867,6 +5946,7 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_persistent_resource( request, @@ -5878,6 +5958,7 @@ def test_delete_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_persistent_resource_rest_bad_request( @@ -6116,11 +6197,15 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_update_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_update_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_update_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.UpdatePersistentResourceRequest.pb( persistent_resource_service.UpdatePersistentResourceRequest() ) @@ -6144,6 +6229,7 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_persistent_resource( request, @@ -6155,6 +6241,7 @@ def test_update_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_reboot_persistent_resource_rest_bad_request( @@ -6240,11 +6327,15 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): transports.PersistentResourceServiceRestInterceptor, "post_reboot_persistent_resource", ) as post, mock.patch.object( + transports.PersistentResourceServiceRestInterceptor, + "post_reboot_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PersistentResourceServiceRestInterceptor, "pre_reboot_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.RebootPersistentResourceRequest.pb( persistent_resource_service.RebootPersistentResourceRequest() ) @@ -6268,6 +6359,7 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.reboot_persistent_resource( request, @@ -6279,6 +6371,7 @@ def test_reboot_persistent_resource_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -7317,11 +7410,15 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_create_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_create_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_create_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.CreatePersistentResourceRequest.pb( persistent_resource_service.CreatePersistentResourceRequest() ) @@ -7345,6 +7442,7 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_persistent_resource( request, @@ -7356,6 +7454,7 @@ async def test_create_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7473,11 +7572,15 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto transports.AsyncPersistentResourceServiceRestInterceptor, "post_get_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_get_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_get_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.GetPersistentResourceRequest.pb( persistent_resource_service.GetPersistentResourceRequest() ) @@ -7503,6 +7606,10 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = persistent_resource.PersistentResource() + post_with_metadata.return_value = ( + persistent_resource.PersistentResource(), + metadata, + ) await client.get_persistent_resource( request, @@ -7514,6 +7621,7 @@ async def test_get_persistent_resource_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7617,11 +7725,15 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep transports.AsyncPersistentResourceServiceRestInterceptor, "post_list_persistent_resources", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_list_persistent_resources_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_list_persistent_resources", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.ListPersistentResourcesRequest.pb( persistent_resource_service.ListPersistentResourcesRequest() ) @@ -7651,6 +7763,10 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep post.return_value = ( persistent_resource_service.ListPersistentResourcesResponse() ) + post_with_metadata.return_value = ( + persistent_resource_service.ListPersistentResourcesResponse(), + metadata, + ) await client.list_persistent_resources( request, @@ -7662,6 +7778,7 @@ async def test_list_persistent_resources_rest_asyncio_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -7763,11 +7880,15 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_delete_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_delete_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_delete_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.DeletePersistentResourceRequest.pb( persistent_resource_service.DeletePersistentResourceRequest() ) @@ -7791,6 +7912,7 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_persistent_resource( request, @@ -7802,6 +7924,7 @@ async def test_delete_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8056,11 +8179,15 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_update_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_update_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_update_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.UpdatePersistentResourceRequest.pb( persistent_resource_service.UpdatePersistentResourceRequest() ) @@ -8084,6 +8211,7 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_persistent_resource( request, @@ -8095,6 +8223,7 @@ async def test_update_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8196,11 +8325,15 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce transports.AsyncPersistentResourceServiceRestInterceptor, "post_reboot_persistent_resource", ) as post, mock.patch.object( + transports.AsyncPersistentResourceServiceRestInterceptor, + "post_reboot_persistent_resource_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPersistentResourceServiceRestInterceptor, "pre_reboot_persistent_resource", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = persistent_resource_service.RebootPersistentResourceRequest.pb( persistent_resource_service.RebootPersistentResourceRequest() ) @@ -8224,6 +8357,7 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.reboot_persistent_resource( request, @@ -8235,6 +8369,7 @@ async def test_reboot_persistent_resource_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_pipeline_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_pipeline_service.py index ff5d304edd..2941bdab07 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_pipeline_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_pipeline_service.py @@ -113,6 +113,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -370,6 +378,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PipelineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PipelineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -9085,6 +9136,7 @@ def test_create_training_pipeline_rest_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -9109,10 +9161,25 @@ def test_create_training_pipeline_rest_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -9326,10 +9393,14 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_create_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_create_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_create_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreateTrainingPipelineRequest.pb( pipeline_service.CreateTrainingPipelineRequest() ) @@ -9355,6 +9426,10 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_training_pipeline.TrainingPipeline() + post_with_metadata.return_value = ( + gca_training_pipeline.TrainingPipeline(), + metadata, + ) client.create_training_pipeline( request, @@ -9366,6 +9441,7 @@ def test_create_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_training_pipeline_rest_bad_request( @@ -9464,10 +9540,14 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_get_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_get_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_get_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetTrainingPipelineRequest.pb( pipeline_service.GetTrainingPipelineRequest() ) @@ -9493,6 +9573,7 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = training_pipeline.TrainingPipeline() + post_with_metadata.return_value = training_pipeline.TrainingPipeline(), metadata client.get_training_pipeline( request, @@ -9504,6 +9585,7 @@ def test_get_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_training_pipelines_rest_bad_request( @@ -9588,10 +9670,14 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_list_training_pipelines" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_list_training_pipelines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_list_training_pipelines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListTrainingPipelinesRequest.pb( pipeline_service.ListTrainingPipelinesRequest() ) @@ -9617,6 +9703,10 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListTrainingPipelinesResponse() + post_with_metadata.return_value = ( + pipeline_service.ListTrainingPipelinesResponse(), + metadata, + ) client.list_training_pipelines( request, @@ -9628,6 +9718,7 @@ def test_list_training_pipelines_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_training_pipeline_rest_bad_request( @@ -9712,10 +9803,14 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_delete_training_pipeline" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_delete_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_delete_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeleteTrainingPipelineRequest.pb( pipeline_service.DeleteTrainingPipelineRequest() ) @@ -9739,6 +9834,7 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_training_pipeline( request, @@ -9750,6 +9846,7 @@ def test_delete_training_pipeline_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_training_pipeline_rest_bad_request( @@ -10159,10 +10256,14 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_create_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_create_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_create_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreatePipelineJobRequest.pb( pipeline_service.CreatePipelineJobRequest() ) @@ -10188,6 +10289,7 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_pipeline_job.PipelineJob() + post_with_metadata.return_value = gca_pipeline_job.PipelineJob(), metadata client.create_pipeline_job( request, @@ -10199,6 +10301,7 @@ def test_create_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_pipeline_job_rest_bad_request( @@ -10305,10 +10408,13 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_get_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, "post_get_pipeline_job_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_get_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetPipelineJobRequest.pb( pipeline_service.GetPipelineJobRequest() ) @@ -10332,6 +10438,7 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_job.PipelineJob() + post_with_metadata.return_value = pipeline_job.PipelineJob(), metadata client.get_pipeline_job( request, @@ -10343,6 +10450,7 @@ def test_get_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_pipeline_jobs_rest_bad_request( @@ -10427,10 +10535,14 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PipelineServiceRestInterceptor, "post_list_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_list_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_list_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListPipelineJobsRequest.pb( pipeline_service.ListPipelineJobsRequest() ) @@ -10456,6 +10568,10 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListPipelineJobsResponse() + post_with_metadata.return_value = ( + pipeline_service.ListPipelineJobsResponse(), + metadata, + ) client.list_pipeline_jobs( request, @@ -10467,6 +10583,7 @@ def test_list_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_pipeline_job_rest_bad_request( @@ -10547,10 +10664,14 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_delete_pipeline_job" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_delete_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_delete_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeletePipelineJobRequest.pb( pipeline_service.DeletePipelineJobRequest() ) @@ -10574,6 +10695,7 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_pipeline_job( request, @@ -10585,6 +10707,7 @@ def test_delete_pipeline_job_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_delete_pipeline_jobs_rest_bad_request( @@ -10665,10 +10788,14 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_batch_delete_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_batch_delete_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_batch_delete_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchDeletePipelineJobsRequest.pb( pipeline_service.BatchDeletePipelineJobsRequest() ) @@ -10692,6 +10819,7 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_delete_pipeline_jobs( request, @@ -10703,6 +10831,7 @@ def test_batch_delete_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_cancel_pipeline_job_rest_bad_request( @@ -10892,10 +11021,14 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): ), mock.patch.object( transports.PipelineServiceRestInterceptor, "post_batch_cancel_pipeline_jobs" ) as post, mock.patch.object( + transports.PipelineServiceRestInterceptor, + "post_batch_cancel_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PipelineServiceRestInterceptor, "pre_batch_cancel_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchCancelPipelineJobsRequest.pb( pipeline_service.BatchCancelPipelineJobsRequest() ) @@ -10919,6 +11052,7 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.batch_cancel_pipeline_jobs( request, @@ -10930,6 +11064,7 @@ def test_batch_cancel_pipeline_jobs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -11956,6 +12091,7 @@ async def test_create_training_pipeline_rest_asyncio_call_success(request_type): "display_name": "display_name_value", "description": "description_value", "version_description": "version_description_value", + "default_checkpoint_id": "default_checkpoint_id_value", "predict_schemata": { "instance_schema_uri": "instance_schema_uri_value", "parameters_schema_uri": "parameters_schema_uri_value", @@ -11980,10 +12116,25 @@ async def test_create_training_pipeline_rest_asyncio_call_success(request_type): "shared_memory_size_mb": 2231, "startup_probe": { "exec_": {"command": ["command_value1", "command_value2"]}, + "http_get": { + "path": "path_value", + "port": 453, + "host": "host_value", + "scheme": "scheme_value", + "http_headers": [ + {"name": "name_value", "value": "value_value"} + ], + }, + "grpc": {"port": 453, "service": "service_value"}, + "tcp_socket": {"port": 453, "host": "host_value"}, "period_seconds": 1489, "timeout_seconds": 1621, + "failure_threshold": 1812, + "success_threshold": 1829, + "initial_delay_seconds": 2214, }, "health_probe": {}, + "liveness_probe": {}, }, "artifact_uri": "artifact_uri_value", "supported_deployment_resources_types": [1], @@ -12204,10 +12355,14 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_create_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_create_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_create_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreateTrainingPipelineRequest.pb( pipeline_service.CreateTrainingPipelineRequest() ) @@ -12233,6 +12388,10 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = gca_training_pipeline.TrainingPipeline() + post_with_metadata.return_value = ( + gca_training_pipeline.TrainingPipeline(), + metadata, + ) await client.create_training_pipeline( request, @@ -12244,6 +12403,7 @@ async def test_create_training_pipeline_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12358,10 +12518,14 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_get_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_get_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_get_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetTrainingPipelineRequest.pb( pipeline_service.GetTrainingPipelineRequest() ) @@ -12387,6 +12551,7 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = training_pipeline.TrainingPipeline() + post_with_metadata.return_value = training_pipeline.TrainingPipeline(), metadata await client.get_training_pipeline( request, @@ -12398,6 +12563,7 @@ async def test_get_training_pipeline_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12498,10 +12664,14 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_list_training_pipelines" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_list_training_pipelines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_list_training_pipelines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListTrainingPipelinesRequest.pb( pipeline_service.ListTrainingPipelinesRequest() ) @@ -12527,6 +12697,10 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = pipeline_service.ListTrainingPipelinesResponse() + post_with_metadata.return_value = ( + pipeline_service.ListTrainingPipelinesResponse(), + metadata, + ) await client.list_training_pipelines( request, @@ -12538,6 +12712,7 @@ async def test_list_training_pipelines_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -12638,10 +12813,14 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept ), mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_delete_training_pipeline" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_delete_training_pipeline_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_delete_training_pipeline" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeleteTrainingPipelineRequest.pb( pipeline_service.DeleteTrainingPipelineRequest() ) @@ -12665,6 +12844,7 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_training_pipeline( request, @@ -12676,6 +12856,7 @@ async def test_delete_training_pipeline_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13117,10 +13298,14 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_create_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_create_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_create_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.CreatePipelineJobRequest.pb( pipeline_service.CreatePipelineJobRequest() ) @@ -13146,6 +13331,7 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_pipeline_job.PipelineJob() + post_with_metadata.return_value = gca_pipeline_job.PipelineJob(), metadata await client.create_pipeline_job( request, @@ -13157,6 +13343,7 @@ async def test_create_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13279,10 +13466,14 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_get_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_get_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_get_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.GetPipelineJobRequest.pb( pipeline_service.GetPipelineJobRequest() ) @@ -13306,6 +13497,7 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_job.PipelineJob() + post_with_metadata.return_value = pipeline_job.PipelineJob(), metadata await client.get_pipeline_job( request, @@ -13317,6 +13509,7 @@ async def test_get_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13417,10 +13610,14 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_list_pipeline_jobs" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_list_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_list_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.ListPipelineJobsRequest.pb( pipeline_service.ListPipelineJobsRequest() ) @@ -13446,6 +13643,10 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = pipeline_service.ListPipelineJobsResponse() + post_with_metadata.return_value = ( + pipeline_service.ListPipelineJobsResponse(), + metadata, + ) await client.list_pipeline_jobs( request, @@ -13457,6 +13658,7 @@ async def test_list_pipeline_jobs_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13553,10 +13755,14 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "post_delete_pipeline_job" ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_delete_pipeline_job_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_delete_pipeline_job" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.DeletePipelineJobRequest.pb( pipeline_service.DeletePipelineJobRequest() ) @@ -13580,6 +13786,7 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_pipeline_job( request, @@ -13591,6 +13798,7 @@ async def test_delete_pipeline_job_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13688,10 +13896,14 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce transports.AsyncPipelineServiceRestInterceptor, "post_batch_delete_pipeline_jobs", ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_batch_delete_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_batch_delete_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchDeletePipelineJobsRequest.pb( pipeline_service.BatchDeletePipelineJobsRequest() ) @@ -13715,6 +13927,7 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_delete_pipeline_jobs( request, @@ -13726,6 +13939,7 @@ async def test_batch_delete_pipeline_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13948,10 +14162,14 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce transports.AsyncPipelineServiceRestInterceptor, "post_batch_cancel_pipeline_jobs", ) as post, mock.patch.object( + transports.AsyncPipelineServiceRestInterceptor, + "post_batch_cancel_pipeline_jobs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPipelineServiceRestInterceptor, "pre_batch_cancel_pipeline_jobs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = pipeline_service.BatchCancelPipelineJobsRequest.pb( pipeline_service.BatchCancelPipelineJobsRequest() ) @@ -13975,6 +14193,7 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.batch_cancel_pipeline_jobs( request, @@ -13986,6 +14205,7 @@ async def test_batch_cancel_pipeline_jobs_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_prediction_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_prediction_service.py index c7bfc0cf8e..2c82819d64 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_prediction_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_prediction_service.py @@ -86,10 +86,19 @@ from google.protobuf import any_pb2 # type: ignore from google.protobuf import duration_pb2 # type: ignore from google.protobuf import struct_pb2 # type: ignore +from google.protobuf import timestamp_pb2 # type: ignore from google.type import latlng_pb2 # type: ignore import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -356,6 +365,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = PredictionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = PredictionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4114,6 +4166,7 @@ def test_generate_content(request_type, transport: str = "grpc"): # Designate an appropriate return value for the call. call.return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) response = client.generate_content(request) @@ -4126,6 +4179,7 @@ def test_generate_content(request_type, transport: str = "grpc"): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" def test_generate_content_non_empty_request_with_auto_populated_field(): @@ -4257,6 +4311,7 @@ async def test_generate_content_async( call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) ) response = await client.generate_content(request) @@ -4270,6 +4325,7 @@ async def test_generate_content_async( # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -7692,6 +7748,7 @@ async def test_generate_content_empty_call_grpc_asyncio(): call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) ) await client.generate_content(request=None) @@ -7849,10 +7906,13 @@ def test_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.PredictRequest.pb( prediction_service.PredictRequest() ) @@ -7878,6 +7938,7 @@ def test_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.PredictResponse() + post_with_metadata.return_value = prediction_service.PredictResponse(), metadata client.predict( request, @@ -7889,6 +7950,7 @@ def test_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_raw_predict_rest_bad_request( @@ -7972,10 +8034,13 @@ def test_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_raw_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.RawPredictRequest.pb( prediction_service.RawPredictRequest() ) @@ -7999,6 +8064,7 @@ def test_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.raw_predict( request, @@ -8010,6 +8076,7 @@ def test_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_raw_predict_rest_bad_request( @@ -8097,10 +8164,14 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_stream_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_stream_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_stream_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamRawPredictRequest.pb( prediction_service.StreamRawPredictRequest() ) @@ -8124,6 +8195,7 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.stream_raw_predict( request, @@ -8135,6 +8207,7 @@ def test_stream_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_direct_predict_rest_bad_request( @@ -8216,10 +8289,13 @@ def test_direct_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_direct_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_direct_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_direct_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectPredictRequest.pb( prediction_service.DirectPredictRequest() ) @@ -8245,6 +8321,10 @@ def test_direct_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectPredictResponse(), + metadata, + ) client.direct_predict( request, @@ -8256,6 +8336,7 @@ def test_direct_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_direct_raw_predict_rest_bad_request( @@ -8340,10 +8421,14 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_direct_raw_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_direct_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_direct_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectRawPredictRequest.pb( prediction_service.DirectRawPredictRequest() ) @@ -8369,6 +8454,10 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectRawPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectRawPredictResponse(), + metadata, + ) client.direct_raw_predict( request, @@ -8380,6 +8469,7 @@ def test_direct_raw_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_direct_predict_rest_error(): @@ -8504,10 +8594,14 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_server_streaming_predict" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_server_streaming_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_server_streaming_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamingPredictRequest.pb( prediction_service.StreamingPredictRequest() ) @@ -8533,6 +8627,10 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.StreamingPredictResponse() + post_with_metadata.return_value = ( + prediction_service.StreamingPredictResponse(), + metadata, + ) client.server_streaming_predict( request, @@ -8544,6 +8642,7 @@ def test_server_streaming_predict_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_streaming_raw_predict_rest_error(): @@ -8639,10 +8738,13 @@ def test_explain_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_explain" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_explain_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_explain" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ExplainRequest.pb( prediction_service.ExplainRequest() ) @@ -8668,6 +8770,7 @@ def test_explain_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.ExplainResponse() + post_with_metadata.return_value = prediction_service.ExplainResponse(), metadata client.explain( request, @@ -8679,6 +8782,7 @@ def test_explain_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_count_tokens_rest_bad_request( @@ -8765,10 +8869,13 @@ def test_count_tokens_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_count_tokens" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, "post_count_tokens_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_count_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.CountTokensRequest.pb( prediction_service.CountTokensRequest() ) @@ -8794,6 +8901,10 @@ def test_count_tokens_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.CountTokensResponse() + post_with_metadata.return_value = ( + prediction_service.CountTokensResponse(), + metadata, + ) client.count_tokens( request, @@ -8805,6 +8916,7 @@ def test_count_tokens_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_generate_content_rest_bad_request( @@ -8853,6 +8965,7 @@ def test_generate_content_rest_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -8870,6 +8983,7 @@ def test_generate_content_rest_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.parametrize("null_interceptor", [True, False]) @@ -8889,10 +9003,14 @@ def test_generate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_generate_content" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -8918,6 +9036,10 @@ def test_generate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) client.generate_content( request, @@ -8929,6 +9051,7 @@ def test_generate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_generate_content_rest_bad_request( @@ -8977,6 +9100,7 @@ def test_stream_generate_content_rest_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -8998,6 +9122,7 @@ def test_stream_generate_content_rest_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.parametrize("null_interceptor", [True, False]) @@ -9017,10 +9142,14 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_stream_generate_content" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_stream_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_stream_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -9046,6 +9175,10 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) client.stream_generate_content( request, @@ -9057,6 +9190,7 @@ def test_stream_generate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_chat_completions_rest_bad_request( @@ -9221,10 +9355,14 @@ def test_chat_completions_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.PredictionServiceRestInterceptor, "post_chat_completions" ) as post, mock.patch.object( + transports.PredictionServiceRestInterceptor, + "post_chat_completions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.PredictionServiceRestInterceptor, "pre_chat_completions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ChatCompletionsRequest.pb( prediction_service.ChatCompletionsRequest() ) @@ -9248,6 +9386,7 @@ def test_chat_completions_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.chat_completions( request, @@ -9259,6 +9398,7 @@ def test_chat_completions_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -10236,10 +10376,13 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, "post_predict_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.PredictRequest.pb( prediction_service.PredictRequest() ) @@ -10265,6 +10408,7 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.PredictResponse() + post_with_metadata.return_value = prediction_service.PredictResponse(), metadata await client.predict( request, @@ -10276,6 +10420,7 @@ async def test_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10375,10 +10520,14 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.RawPredictRequest.pb( prediction_service.RawPredictRequest() ) @@ -10402,6 +10551,7 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.raw_predict( request, @@ -10413,6 +10563,7 @@ async def test_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10514,10 +10665,14 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_stream_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_stream_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_stream_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamRawPredictRequest.pb( prediction_service.StreamRawPredictRequest() ) @@ -10541,6 +10696,7 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.stream_raw_predict( request, @@ -10552,6 +10708,7 @@ async def test_stream_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10649,10 +10806,14 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_direct_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_direct_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_direct_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectPredictRequest.pb( prediction_service.DirectPredictRequest() ) @@ -10678,6 +10839,10 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectPredictResponse(), + metadata, + ) await client.direct_predict( request, @@ -10689,6 +10854,7 @@ async def test_direct_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10789,10 +10955,14 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_direct_raw_predict" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_direct_raw_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_direct_raw_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.DirectRawPredictRequest.pb( prediction_service.DirectRawPredictRequest() ) @@ -10818,6 +10988,10 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.DirectRawPredictResponse() + post_with_metadata.return_value = ( + prediction_service.DirectRawPredictResponse(), + metadata, + ) await client.direct_raw_predict( request, @@ -10829,6 +11003,7 @@ async def test_direct_raw_predict_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10983,10 +11158,14 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept transports.AsyncPredictionServiceRestInterceptor, "post_server_streaming_predict", ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_server_streaming_predict_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_server_streaming_predict" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.StreamingPredictRequest.pb( prediction_service.StreamingPredictRequest() ) @@ -11012,6 +11191,10 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept ] pre.return_value = request, metadata post.return_value = prediction_service.StreamingPredictResponse() + post_with_metadata.return_value = ( + prediction_service.StreamingPredictResponse(), + metadata, + ) await client.server_streaming_predict( request, @@ -11023,6 +11206,7 @@ async def test_server_streaming_predict_rest_asyncio_interceptors(null_intercept pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11141,10 +11325,13 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_explain" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, "post_explain_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_explain" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ExplainRequest.pb( prediction_service.ExplainRequest() ) @@ -11170,6 +11357,7 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.ExplainResponse() + post_with_metadata.return_value = prediction_service.ExplainResponse(), metadata await client.explain( request, @@ -11181,6 +11369,7 @@ async def test_explain_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11283,10 +11472,14 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_count_tokens" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_count_tokens_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_count_tokens" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.CountTokensRequest.pb( prediction_service.CountTokensRequest() ) @@ -11312,6 +11505,10 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.CountTokensResponse() + post_with_metadata.return_value = ( + prediction_service.CountTokensResponse(), + metadata, + ) await client.count_tokens( request, @@ -11323,6 +11520,7 @@ async def test_count_tokens_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11380,6 +11578,7 @@ async def test_generate_content_rest_asyncio_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -11399,6 +11598,7 @@ async def test_generate_content_rest_asyncio_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -11423,10 +11623,14 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_generate_content" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -11452,6 +11656,10 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) await client.generate_content( request, @@ -11463,6 +11671,7 @@ async def test_generate_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11520,6 +11729,7 @@ async def test_stream_generate_content_rest_asyncio_call_success(request_type): # Designate an appropriate value for the returned response. return_value = prediction_service.GenerateContentResponse( model_version="model_version_value", + response_id="response_id_value", ) # Wrap the value into a proper Response obj @@ -11541,6 +11751,7 @@ async def test_stream_generate_content_rest_asyncio_call_success(request_type): # Establish that the response is the type that we expect. assert isinstance(response, prediction_service.GenerateContentResponse) assert response.model_version == "model_version_value" + assert response.response_id == "response_id_value" @pytest.mark.asyncio @@ -11565,10 +11776,14 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_stream_generate_content" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_stream_generate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_stream_generate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.GenerateContentRequest.pb( prediction_service.GenerateContentRequest() ) @@ -11594,6 +11809,10 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = prediction_service.GenerateContentResponse() + post_with_metadata.return_value = ( + prediction_service.GenerateContentResponse(), + metadata, + ) await client.stream_generate_content( request, @@ -11605,6 +11824,7 @@ async def test_stream_generate_content_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11783,10 +12003,14 @@ async def test_chat_completions_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "post_chat_completions" ) as post, mock.patch.object( + transports.AsyncPredictionServiceRestInterceptor, + "post_chat_completions_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncPredictionServiceRestInterceptor, "pre_chat_completions" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = prediction_service.ChatCompletionsRequest.pb( prediction_service.ChatCompletionsRequest() ) @@ -11810,6 +12034,7 @@ async def test_chat_completions_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.chat_completions( request, @@ -11821,6 +12046,7 @@ async def test_chat_completions_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_execution_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_execution_service.py index 855b569109..d5da1a55a4 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_execution_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_execution_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -377,6 +385,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ReasoningEngineExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ReasoningEngineExecutionServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -2278,11 +2329,15 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): transports.ReasoningEngineExecutionServiceRestInterceptor, "post_query_reasoning_engine", ) as post, mock.patch.object( + transports.ReasoningEngineExecutionServiceRestInterceptor, + "post_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineExecutionServiceRestInterceptor, "pre_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_execution_service.QueryReasoningEngineRequest.pb( reasoning_engine_execution_service.QueryReasoningEngineRequest() ) @@ -2312,6 +2367,10 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): post.return_value = ( reasoning_engine_execution_service.QueryReasoningEngineResponse() ) + post_with_metadata.return_value = ( + reasoning_engine_execution_service.QueryReasoningEngineResponse(), + metadata, + ) client.query_reasoning_engine( request, @@ -2323,6 +2382,7 @@ def test_query_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stream_query_reasoning_engine_rest_bad_request( @@ -2415,11 +2475,15 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): transports.ReasoningEngineExecutionServiceRestInterceptor, "post_stream_query_reasoning_engine", ) as post, mock.patch.object( + transports.ReasoningEngineExecutionServiceRestInterceptor, + "post_stream_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineExecutionServiceRestInterceptor, "pre_stream_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest.pb( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest() @@ -2445,6 +2509,7 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata client.stream_query_reasoning_engine( request, @@ -2456,6 +2521,7 @@ def test_stream_query_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -3253,11 +3319,15 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "post_query_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineExecutionServiceRestInterceptor, + "post_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "pre_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_execution_service.QueryReasoningEngineRequest.pb( reasoning_engine_execution_service.QueryReasoningEngineRequest() ) @@ -3287,6 +3357,10 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor post.return_value = ( reasoning_engine_execution_service.QueryReasoningEngineResponse() ) + post_with_metadata.return_value = ( + reasoning_engine_execution_service.QueryReasoningEngineResponse(), + metadata, + ) await client.query_reasoning_engine( request, @@ -3298,6 +3372,7 @@ async def test_query_reasoning_engine_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -3406,11 +3481,15 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "post_stream_query_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineExecutionServiceRestInterceptor, + "post_stream_query_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineExecutionServiceRestInterceptor, "pre_stream_query_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = ( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest.pb( reasoning_engine_execution_service.StreamQueryReasoningEngineRequest() @@ -3436,6 +3515,7 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = httpbody_pb2.HttpBody() + post_with_metadata.return_value = httpbody_pb2.HttpBody(), metadata await client.stream_query_reasoning_engine( request, @@ -3447,6 +3527,7 @@ async def test_stream_query_reasoning_engine_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_service.py index bd7a704224..b1f7271527 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_reasoning_engine_service.py @@ -93,6 +93,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -367,6 +375,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ReasoningEngineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ReasoningEngineServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4679,10 +4730,14 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_create_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_create_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_create_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.CreateReasoningEngineRequest.pb( reasoning_engine_service.CreateReasoningEngineRequest() ) @@ -4706,6 +4761,7 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_reasoning_engine( request, @@ -4717,6 +4773,7 @@ def test_create_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_reasoning_engine_rest_bad_request( @@ -4811,10 +4868,14 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_get_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_get_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_get_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.GetReasoningEngineRequest.pb( reasoning_engine_service.GetReasoningEngineRequest() ) @@ -4840,6 +4901,7 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine.ReasoningEngine() + post_with_metadata.return_value = reasoning_engine.ReasoningEngine(), metadata client.get_reasoning_engine( request, @@ -4851,6 +4913,7 @@ def test_get_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_reasoning_engines_rest_bad_request( @@ -4937,10 +5000,14 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_list_reasoning_engines" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_list_reasoning_engines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_list_reasoning_engines" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.ListReasoningEnginesRequest.pb( reasoning_engine_service.ListReasoningEnginesRequest() ) @@ -4966,6 +5033,10 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine_service.ListReasoningEnginesResponse() + post_with_metadata.return_value = ( + reasoning_engine_service.ListReasoningEnginesResponse(), + metadata, + ) client.list_reasoning_engines( request, @@ -4977,6 +5048,7 @@ def test_list_reasoning_engines_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_reasoning_engine_rest_bad_request( @@ -5151,10 +5223,14 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_update_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_update_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_update_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.UpdateReasoningEngineRequest.pb( reasoning_engine_service.UpdateReasoningEngineRequest() ) @@ -5178,6 +5254,7 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_reasoning_engine( request, @@ -5189,6 +5266,7 @@ def test_update_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_reasoning_engine_rest_bad_request( @@ -5273,10 +5351,14 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "post_delete_reasoning_engine" ) as post, mock.patch.object( + transports.ReasoningEngineServiceRestInterceptor, + "post_delete_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.ReasoningEngineServiceRestInterceptor, "pre_delete_reasoning_engine" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.DeleteReasoningEngineRequest.pb( reasoning_engine_service.DeleteReasoningEngineRequest() ) @@ -5300,6 +5382,7 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_reasoning_engine( request, @@ -5311,6 +5394,7 @@ def test_delete_reasoning_engine_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6264,11 +6348,15 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_create_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_create_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_create_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.CreateReasoningEngineRequest.pb( reasoning_engine_service.CreateReasoningEngineRequest() ) @@ -6292,6 +6380,7 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_reasoning_engine( request, @@ -6303,6 +6392,7 @@ async def test_create_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6414,11 +6504,15 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): transports.AsyncReasoningEngineServiceRestInterceptor, "post_get_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_get_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_get_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.GetReasoningEngineRequest.pb( reasoning_engine_service.GetReasoningEngineRequest() ) @@ -6444,6 +6538,7 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = reasoning_engine.ReasoningEngine() + post_with_metadata.return_value = reasoning_engine.ReasoningEngine(), metadata await client.get_reasoning_engine( request, @@ -6455,6 +6550,7 @@ async def test_get_reasoning_engine_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6558,11 +6654,15 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor transports.AsyncReasoningEngineServiceRestInterceptor, "post_list_reasoning_engines", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_list_reasoning_engines_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_list_reasoning_engines", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.ListReasoningEnginesRequest.pb( reasoning_engine_service.ListReasoningEnginesRequest() ) @@ -6588,6 +6688,10 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = reasoning_engine_service.ListReasoningEnginesResponse() + post_with_metadata.return_value = ( + reasoning_engine_service.ListReasoningEnginesResponse(), + metadata, + ) await client.list_reasoning_engines( request, @@ -6599,6 +6703,7 @@ async def test_list_reasoning_engines_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6790,11 +6895,15 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_update_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_update_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_update_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.UpdateReasoningEngineRequest.pb( reasoning_engine_service.UpdateReasoningEngineRequest() ) @@ -6818,6 +6927,7 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_reasoning_engine( request, @@ -6829,6 +6939,7 @@ async def test_update_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6930,11 +7041,15 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto transports.AsyncReasoningEngineServiceRestInterceptor, "post_delete_reasoning_engine", ) as post, mock.patch.object( + transports.AsyncReasoningEngineServiceRestInterceptor, + "post_delete_reasoning_engine_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncReasoningEngineServiceRestInterceptor, "pre_delete_reasoning_engine", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = reasoning_engine_service.DeleteReasoningEngineRequest.pb( reasoning_engine_service.DeleteReasoningEngineRequest() ) @@ -6958,6 +7073,7 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_reasoning_engine( request, @@ -6969,6 +7085,7 @@ async def test_delete_reasoning_engine_rest_asyncio_interceptors(null_intercepto pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_schedule_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_schedule_service.py index 44dcd27823..f649f3a305 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_schedule_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_schedule_service.py @@ -120,6 +120,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -377,6 +385,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = ScheduleServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = ScheduleServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -6032,10 +6083,13 @@ def test_create_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_create_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_create_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_create_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.CreateScheduleRequest.pb( schedule_service.CreateScheduleRequest() ) @@ -6059,6 +6113,7 @@ def test_create_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata client.create_schedule( request, @@ -6070,6 +6125,7 @@ def test_create_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_schedule_rest_bad_request( @@ -6150,10 +6206,13 @@ def test_delete_schedule_rest_interceptors(null_interceptor): ), mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_delete_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_delete_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_delete_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.DeleteScheduleRequest.pb( schedule_service.DeleteScheduleRequest() ) @@ -6177,6 +6236,7 @@ def test_delete_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_schedule( request, @@ -6188,6 +6248,7 @@ def test_delete_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_schedule_rest_bad_request( @@ -6287,10 +6348,13 @@ def test_get_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_get_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_get_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_get_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.GetScheduleRequest.pb( schedule_service.GetScheduleRequest() ) @@ -6314,6 +6378,7 @@ def test_get_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule.Schedule() + post_with_metadata.return_value = schedule.Schedule(), metadata client.get_schedule( request, @@ -6325,6 +6390,7 @@ def test_get_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_schedules_rest_bad_request( @@ -6409,10 +6475,13 @@ def test_list_schedules_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_list_schedules" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_list_schedules_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_list_schedules" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.ListSchedulesRequest.pb( schedule_service.ListSchedulesRequest() ) @@ -6438,6 +6507,10 @@ def test_list_schedules_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule_service.ListSchedulesResponse() + post_with_metadata.return_value = ( + schedule_service.ListSchedulesResponse(), + metadata, + ) client.list_schedules( request, @@ -6449,6 +6522,7 @@ def test_list_schedules_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_pause_schedule_rest_bad_request( @@ -7188,10 +7262,13 @@ def test_update_schedule_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.ScheduleServiceRestInterceptor, "post_update_schedule" ) as post, mock.patch.object( + transports.ScheduleServiceRestInterceptor, "post_update_schedule_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.ScheduleServiceRestInterceptor, "pre_update_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.UpdateScheduleRequest.pb( schedule_service.UpdateScheduleRequest() ) @@ -7215,6 +7292,7 @@ def test_update_schedule_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata client.update_schedule( request, @@ -7226,6 +7304,7 @@ def test_update_schedule_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -8559,10 +8638,14 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_create_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_create_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_create_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.CreateScheduleRequest.pb( schedule_service.CreateScheduleRequest() ) @@ -8586,6 +8669,7 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata await client.create_schedule( request, @@ -8597,6 +8681,7 @@ async def test_create_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8693,10 +8778,14 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_delete_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_delete_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_delete_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.DeleteScheduleRequest.pb( schedule_service.DeleteScheduleRequest() ) @@ -8720,6 +8809,7 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_schedule( request, @@ -8731,6 +8821,7 @@ async def test_delete_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8846,10 +8937,14 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_get_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_get_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_get_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.GetScheduleRequest.pb( schedule_service.GetScheduleRequest() ) @@ -8873,6 +8968,7 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule.Schedule() + post_with_metadata.return_value = schedule.Schedule(), metadata await client.get_schedule( request, @@ -8884,6 +8980,7 @@ async def test_get_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -8984,10 +9081,14 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_list_schedules" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_list_schedules_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_list_schedules" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.ListSchedulesRequest.pb( schedule_service.ListSchedulesRequest() ) @@ -9013,6 +9114,10 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = schedule_service.ListSchedulesResponse() + post_with_metadata.return_value = ( + schedule_service.ListSchedulesResponse(), + metadata, + ) await client.list_schedules( request, @@ -9024,6 +9129,7 @@ async def test_list_schedules_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -9811,10 +9917,14 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "post_update_schedule" ) as post, mock.patch.object( + transports.AsyncScheduleServiceRestInterceptor, + "post_update_schedule_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncScheduleServiceRestInterceptor, "pre_update_schedule" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = schedule_service.UpdateScheduleRequest.pb( schedule_service.UpdateScheduleRequest() ) @@ -9838,6 +9948,7 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_schedule.Schedule() + post_with_metadata.return_value = gca_schedule.Schedule(), metadata await client.update_schedule( request, @@ -9849,6 +9960,7 @@ async def test_update_schedule_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_specialist_pool_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_specialist_pool_service.py index 187cc99ba0..445e3ac331 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_specialist_pool_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_specialist_pool_service.py @@ -89,6 +89,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -363,6 +371,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = SpecialistPoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = SpecialistPoolServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -4691,10 +4742,14 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_create_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_create_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_create_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.CreateSpecialistPoolRequest.pb( specialist_pool_service.CreateSpecialistPoolRequest() ) @@ -4718,6 +4773,7 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_specialist_pool( request, @@ -4729,6 +4785,7 @@ def test_create_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_specialist_pool_rest_bad_request( @@ -4827,10 +4884,14 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_get_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_get_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_get_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.GetSpecialistPoolRequest.pb( specialist_pool_service.GetSpecialistPoolRequest() ) @@ -4856,6 +4917,7 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool.SpecialistPool() + post_with_metadata.return_value = specialist_pool.SpecialistPool(), metadata client.get_specialist_pool( request, @@ -4867,6 +4929,7 @@ def test_get_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_specialist_pools_rest_bad_request( @@ -4953,10 +5016,14 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_list_specialist_pools" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_list_specialist_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_list_specialist_pools" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.ListSpecialistPoolsRequest.pb( specialist_pool_service.ListSpecialistPoolsRequest() ) @@ -4982,6 +5049,10 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool_service.ListSpecialistPoolsResponse() + post_with_metadata.return_value = ( + specialist_pool_service.ListSpecialistPoolsResponse(), + metadata, + ) client.list_specialist_pools( request, @@ -4993,6 +5064,7 @@ def test_list_specialist_pools_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_specialist_pool_rest_bad_request( @@ -5077,10 +5149,14 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_delete_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_delete_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_delete_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.DeleteSpecialistPoolRequest.pb( specialist_pool_service.DeleteSpecialistPoolRequest() ) @@ -5104,6 +5180,7 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_specialist_pool( request, @@ -5115,6 +5192,7 @@ def test_delete_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_specialist_pool_rest_bad_request( @@ -5289,10 +5367,14 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): ), mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "post_update_specialist_pool" ) as post, mock.patch.object( + transports.SpecialistPoolServiceRestInterceptor, + "post_update_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.SpecialistPoolServiceRestInterceptor, "pre_update_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.UpdateSpecialistPoolRequest.pb( specialist_pool_service.UpdateSpecialistPoolRequest() ) @@ -5316,6 +5398,7 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_specialist_pool( request, @@ -5327,6 +5410,7 @@ def test_update_specialist_pool_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -6280,11 +6364,15 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_create_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_create_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_create_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.CreateSpecialistPoolRequest.pb( specialist_pool_service.CreateSpecialistPoolRequest() ) @@ -6308,6 +6396,7 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_specialist_pool( request, @@ -6319,6 +6408,7 @@ async def test_create_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6433,10 +6523,14 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "post_get_specialist_pool" ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_get_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_get_specialist_pool" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.GetSpecialistPoolRequest.pb( specialist_pool_service.GetSpecialistPoolRequest() ) @@ -6462,6 +6556,7 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = specialist_pool.SpecialistPool() + post_with_metadata.return_value = specialist_pool.SpecialistPool(), metadata await client.get_specialist_pool( request, @@ -6473,6 +6568,7 @@ async def test_get_specialist_pool_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6576,11 +6672,15 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) transports.AsyncSpecialistPoolServiceRestInterceptor, "post_list_specialist_pools", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_list_specialist_pools_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_list_specialist_pools", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.ListSpecialistPoolsRequest.pb( specialist_pool_service.ListSpecialistPoolsRequest() ) @@ -6606,6 +6706,10 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = specialist_pool_service.ListSpecialistPoolsResponse() + post_with_metadata.return_value = ( + specialist_pool_service.ListSpecialistPoolsResponse(), + metadata, + ) await client.list_specialist_pools( request, @@ -6617,6 +6721,7 @@ async def test_list_specialist_pools_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6718,11 +6823,15 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_delete_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_delete_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_delete_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.DeleteSpecialistPoolRequest.pb( specialist_pool_service.DeleteSpecialistPoolRequest() ) @@ -6746,6 +6855,7 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_specialist_pool( request, @@ -6757,6 +6867,7 @@ async def test_delete_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -6948,11 +7059,15 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor transports.AsyncSpecialistPoolServiceRestInterceptor, "post_update_specialist_pool", ) as post, mock.patch.object( + transports.AsyncSpecialistPoolServiceRestInterceptor, + "post_update_specialist_pool_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncSpecialistPoolServiceRestInterceptor, "pre_update_specialist_pool", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = specialist_pool_service.UpdateSpecialistPoolRequest.pb( specialist_pool_service.UpdateSpecialistPoolRequest() ) @@ -6976,6 +7091,7 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_specialist_pool( request, @@ -6987,6 +7103,7 @@ async def test_update_specialist_pool_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_tensorboard_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_tensorboard_service.py index 1694fd457b..8eeb8619a8 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_tensorboard_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_tensorboard_service.py @@ -102,6 +102,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -372,6 +380,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = TensorboardServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = TensorboardServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -21132,10 +21183,14 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_create_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRequest.pb( tensorboard_service.CreateTensorboardRequest() ) @@ -21159,6 +21214,7 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_tensorboard( request, @@ -21170,6 +21226,7 @@ def test_create_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_rest_bad_request( @@ -21270,10 +21327,14 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRequest.pb( tensorboard_service.GetTensorboardRequest() ) @@ -21297,6 +21358,7 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard.Tensorboard() + post_with_metadata.return_value = tensorboard.Tensorboard(), metadata client.get_tensorboard( request, @@ -21308,6 +21370,7 @@ def test_get_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_rest_bad_request( @@ -21478,10 +21541,14 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_update_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRequest.pb( tensorboard_service.UpdateTensorboardRequest() ) @@ -21505,6 +21572,7 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_tensorboard( request, @@ -21516,6 +21584,7 @@ def test_update_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboards_rest_bad_request( @@ -21600,10 +21669,14 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_list_tensorboards" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboards_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboards" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardsRequest.pb( tensorboard_service.ListTensorboardsRequest() ) @@ -21629,6 +21702,10 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardsResponse(), + metadata, + ) client.list_tensorboards( request, @@ -21640,6 +21717,7 @@ def test_list_tensorboards_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_rest_bad_request( @@ -21720,10 +21798,14 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRequest.pb( tensorboard_service.DeleteTensorboardRequest() ) @@ -21747,6 +21829,7 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard( request, @@ -21758,6 +21841,7 @@ def test_delete_tensorboard_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_usage_rest_bad_request( @@ -21843,10 +21927,14 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_usage" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_usage_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_usage" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardUsageRequest.pb( tensorboard_service.ReadTensorboardUsageRequest() ) @@ -21872,6 +21960,10 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardUsageResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardUsageResponse(), + metadata, + ) client.read_tensorboard_usage( request, @@ -21883,6 +21975,7 @@ def test_read_tensorboard_usage_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_size_rest_bad_request( @@ -21971,10 +22064,14 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_size" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_size_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_size" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardSizeRequest.pb( tensorboard_service.ReadTensorboardSizeRequest() ) @@ -22000,6 +22097,10 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardSizeResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardSizeResponse(), + metadata, + ) client.read_tensorboard_size( request, @@ -22011,6 +22112,7 @@ def test_read_tensorboard_size_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_experiment_rest_bad_request( @@ -22185,11 +22287,15 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardExperimentRequest.pb( tensorboard_service.CreateTensorboardExperimentRequest() ) @@ -22215,6 +22321,10 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.create_tensorboard_experiment( request, @@ -22226,6 +22336,7 @@ def test_create_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_experiment_rest_bad_request( @@ -22322,10 +22433,14 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_experiment" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_experiment" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardExperimentRequest.pb( tensorboard_service.GetTensorboardExperimentRequest() ) @@ -22351,6 +22466,10 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.get_tensorboard_experiment( request, @@ -22362,6 +22481,7 @@ def test_get_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_experiment_rest_bad_request( @@ -22544,11 +22664,15 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardExperimentRequest.pb( tensorboard_service.UpdateTensorboardExperimentRequest() ) @@ -22574,6 +22698,10 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) client.update_tensorboard_experiment( request, @@ -22585,6 +22713,7 @@ def test_update_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_experiments_rest_bad_request( @@ -22672,10 +22801,14 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_experiments", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_experiments_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_experiments" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardExperimentsRequest.pb( tensorboard_service.ListTensorboardExperimentsRequest() ) @@ -22701,6 +22834,10 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardExperimentsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardExperimentsResponse(), + metadata, + ) client.list_tensorboard_experiments( request, @@ -22712,6 +22849,7 @@ def test_list_tensorboard_experiments_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_experiment_rest_bad_request( @@ -22797,11 +22935,15 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_experiment", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardExperimentRequest.pb( tensorboard_service.DeleteTensorboardExperimentRequest() ) @@ -22825,6 +22967,7 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_experiment( request, @@ -22836,6 +22979,7 @@ def test_delete_tensorboard_experiment_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_run_rest_bad_request( @@ -23008,10 +23152,14 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRunRequest.pb( tensorboard_service.CreateTensorboardRunRequest() ) @@ -23037,6 +23185,7 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata client.create_tensorboard_run( request, @@ -23048,6 +23197,7 @@ def test_create_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_tensorboard_runs_rest_bad_request( @@ -23136,11 +23286,15 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_batch_create_tensorboard_runs", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_runs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardRunsRequest.pb( tensorboard_service.BatchCreateTensorboardRunsRequest() ) @@ -23166,6 +23320,10 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.BatchCreateTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardRunsResponse(), + metadata, + ) client.batch_create_tensorboard_runs( request, @@ -23177,6 +23335,7 @@ def test_batch_create_tensorboard_runs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_run_rest_bad_request( @@ -23271,10 +23430,14 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRunRequest.pb( tensorboard_service.GetTensorboardRunRequest() ) @@ -23300,6 +23463,7 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_run.TensorboardRun() + post_with_metadata.return_value = tensorboard_run.TensorboardRun(), metadata client.get_tensorboard_run( request, @@ -23311,6 +23475,7 @@ def test_get_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_run_rest_bad_request( @@ -23487,10 +23652,14 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRunRequest.pb( tensorboard_service.UpdateTensorboardRunRequest() ) @@ -23516,6 +23685,7 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata client.update_tensorboard_run( request, @@ -23527,6 +23697,7 @@ def test_update_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_runs_rest_bad_request( @@ -23615,10 +23786,14 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_runs" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_runs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardRunsRequest.pb( tensorboard_service.ListTensorboardRunsRequest() ) @@ -23644,6 +23819,10 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardRunsResponse(), + metadata, + ) client.list_tensorboard_runs( request, @@ -23655,6 +23834,7 @@ def test_list_tensorboard_runs_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_run_rest_bad_request( @@ -23739,10 +23919,14 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): ), mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_run" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRunRequest.pb( tensorboard_service.DeleteTensorboardRunRequest() ) @@ -23766,6 +23950,7 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_run( request, @@ -23777,6 +23962,7 @@ def test_delete_tensorboard_run_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_create_tensorboard_time_series_rest_bad_request( @@ -23867,11 +24053,15 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor transports.TensorboardServiceRestInterceptor, "post_batch_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardTimeSeriesRequest.pb( tensorboard_service.BatchCreateTensorboardTimeSeriesRequest() ) @@ -23901,6 +24091,10 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor post.return_value = ( tensorboard_service.BatchCreateTensorboardTimeSeriesResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse(), + metadata, + ) client.batch_create_tensorboard_time_series( request, @@ -23912,6 +24106,7 @@ def test_batch_create_tensorboard_time_series_rest_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_tensorboard_time_series_rest_bad_request( @@ -24105,11 +24300,15 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardTimeSeriesRequest.pb( tensorboard_service.CreateTensorboardTimeSeriesRequest() ) @@ -24135,6 +24334,10 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.create_tensorboard_time_series( request, @@ -24146,6 +24349,7 @@ def test_create_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_tensorboard_time_series_rest_bad_request( @@ -24249,10 +24453,14 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_get_tensorboard_time_series" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_get_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_get_tensorboard_time_series" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardTimeSeriesRequest.pb( tensorboard_service.GetTensorboardTimeSeriesRequest() ) @@ -24278,6 +24486,10 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.get_tensorboard_time_series( request, @@ -24289,6 +24501,7 @@ def test_get_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_tensorboard_time_series_rest_bad_request( @@ -24486,11 +24699,15 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_update_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_update_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_update_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardTimeSeriesRequest.pb( tensorboard_service.UpdateTensorboardTimeSeriesRequest() ) @@ -24516,6 +24733,10 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) client.update_tensorboard_time_series( request, @@ -24527,6 +24748,7 @@ def test_update_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_tensorboard_time_series_rest_bad_request( @@ -24618,10 +24840,14 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_list_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_list_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_list_tensorboard_time_series" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardTimeSeriesRequest.pb( tensorboard_service.ListTensorboardTimeSeriesRequest() ) @@ -24647,6 +24873,10 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardTimeSeriesResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardTimeSeriesResponse(), + metadata, + ) client.list_tensorboard_time_series( request, @@ -24658,6 +24888,7 @@ def test_list_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_tensorboard_time_series_rest_bad_request( @@ -24743,11 +24974,15 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_delete_tensorboard_time_series", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_delete_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_delete_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardTimeSeriesRequest.pb( tensorboard_service.DeleteTensorboardTimeSeriesRequest() ) @@ -24771,6 +25006,7 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_tensorboard_time_series( request, @@ -24782,6 +25018,7 @@ def test_delete_tensorboard_time_series_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_batch_read_tensorboard_time_series_data_rest_bad_request( @@ -24874,11 +25111,15 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep transports.TensorboardServiceRestInterceptor, "post_batch_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_batch_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_batch_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest() ) @@ -24908,6 +25149,10 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep post.return_value = ( tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse(), + metadata, + ) client.batch_read_tensorboard_time_series_data( request, @@ -24919,6 +25164,7 @@ def test_batch_read_tensorboard_time_series_data_rest_interceptors(null_intercep pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_time_series_data_rest_bad_request( @@ -25009,11 +25255,15 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ReadTensorboardTimeSeriesDataRequest() ) @@ -25041,6 +25291,10 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardTimeSeriesDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardTimeSeriesDataResponse(), + metadata, + ) client.read_tensorboard_time_series_data( request, @@ -25052,6 +25306,7 @@ def test_read_tensorboard_time_series_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_read_tensorboard_blob_data_rest_bad_request( @@ -25143,10 +25398,14 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_read_tensorboard_blob_data" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_read_tensorboard_blob_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_read_tensorboard_blob_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardBlobDataRequest.pb( tensorboard_service.ReadTensorboardBlobDataRequest() ) @@ -25172,6 +25431,10 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardBlobDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardBlobDataResponse(), + metadata, + ) client.read_tensorboard_blob_data( request, @@ -25183,6 +25446,7 @@ def test_read_tensorboard_blob_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_tensorboard_experiment_data_rest_bad_request( @@ -25273,11 +25537,15 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): transports.TensorboardServiceRestInterceptor, "post_write_tensorboard_experiment_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_write_tensorboard_experiment_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_write_tensorboard_experiment_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardExperimentDataRequest.pb( tensorboard_service.WriteTensorboardExperimentDataRequest() ) @@ -25305,6 +25573,10 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardExperimentDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardExperimentDataResponse(), + metadata, + ) client.write_tensorboard_experiment_data( request, @@ -25316,6 +25588,7 @@ def test_write_tensorboard_experiment_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_write_tensorboard_run_data_rest_bad_request( @@ -25403,10 +25676,14 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.TensorboardServiceRestInterceptor, "post_write_tensorboard_run_data" ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_write_tensorboard_run_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_write_tensorboard_run_data" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardRunDataRequest.pb( tensorboard_service.WriteTensorboardRunDataRequest() ) @@ -25432,6 +25709,10 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardRunDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardRunDataResponse(), + metadata, + ) client.write_tensorboard_run_data( request, @@ -25443,6 +25724,7 @@ def test_write_tensorboard_run_data_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_export_tensorboard_time_series_data_rest_bad_request( @@ -25534,11 +25816,15 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) transports.TensorboardServiceRestInterceptor, "post_export_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.TensorboardServiceRestInterceptor, + "post_export_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.TensorboardServiceRestInterceptor, "pre_export_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ExportTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ExportTensorboardTimeSeriesDataRequest() ) @@ -25568,6 +25854,10 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) post.return_value = ( tensorboard_service.ExportTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.ExportTensorboardTimeSeriesDataResponse(), + metadata, + ) client.export_tensorboard_time_series_data( request, @@ -25579,6 +25869,7 @@ def test_export_tensorboard_time_series_data_rest_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -27075,10 +27366,14 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRequest.pb( tensorboard_service.CreateTensorboardRequest() ) @@ -27102,6 +27397,7 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_tensorboard( request, @@ -27113,6 +27409,7 @@ async def test_create_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27229,10 +27526,14 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRequest.pb( tensorboard_service.GetTensorboardRequest() ) @@ -27256,6 +27557,7 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard.Tensorboard() + post_with_metadata.return_value = tensorboard.Tensorboard(), metadata await client.get_tensorboard( request, @@ -27267,6 +27569,7 @@ async def test_get_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27453,10 +27756,14 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRequest.pb( tensorboard_service.UpdateTensorboardRequest() ) @@ -27480,6 +27787,7 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_tensorboard( request, @@ -27491,6 +27799,7 @@ async def test_update_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27591,10 +27900,14 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboards" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboards_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboards" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardsRequest.pb( tensorboard_service.ListTensorboardsRequest() ) @@ -27620,6 +27933,10 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardsResponse(), + metadata, + ) await client.list_tensorboards( request, @@ -27631,6 +27948,7 @@ async def test_list_tensorboards_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27727,10 +28045,14 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRequest.pb( tensorboard_service.DeleteTensorboardRequest() ) @@ -27754,6 +28076,7 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard( request, @@ -27765,6 +28088,7 @@ async def test_delete_tensorboard_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -27866,10 +28190,14 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_usage" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_usage_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_usage" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardUsageRequest.pb( tensorboard_service.ReadTensorboardUsageRequest() ) @@ -27895,6 +28223,10 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardUsageResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardUsageResponse(), + metadata, + ) await client.read_tensorboard_usage( request, @@ -27906,6 +28238,7 @@ async def test_read_tensorboard_usage_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28010,10 +28343,14 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_size" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_size_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_size" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardSizeRequest.pb( tensorboard_service.ReadTensorboardSizeRequest() ) @@ -28039,6 +28376,10 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardSizeResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardSizeResponse(), + metadata, + ) await client.read_tensorboard_size( request, @@ -28050,6 +28391,7 @@ async def test_read_tensorboard_size_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28242,11 +28584,15 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardExperimentRequest.pb( tensorboard_service.CreateTensorboardExperimentRequest() ) @@ -28272,6 +28618,10 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.create_tensorboard_experiment( request, @@ -28283,6 +28633,7 @@ async def test_create_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28396,11 +28747,15 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardExperimentRequest.pb( tensorboard_service.GetTensorboardExperimentRequest() ) @@ -28426,6 +28781,10 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.get_tensorboard_experiment( request, @@ -28437,6 +28796,7 @@ async def test_get_tensorboard_experiment_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28637,11 +28997,15 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardExperimentRequest.pb( tensorboard_service.UpdateTensorboardExperimentRequest() ) @@ -28667,6 +29031,10 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_experiment.TensorboardExperiment() + post_with_metadata.return_value = ( + gca_tensorboard_experiment.TensorboardExperiment(), + metadata, + ) await client.update_tensorboard_experiment( request, @@ -28678,6 +29046,7 @@ async def test_update_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28781,11 +29150,15 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_experiments", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_experiments_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_experiments", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardExperimentsRequest.pb( tensorboard_service.ListTensorboardExperimentsRequest() ) @@ -28811,6 +29184,10 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardExperimentsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardExperimentsResponse(), + metadata, + ) await client.list_tensorboard_experiments( request, @@ -28822,6 +29199,7 @@ async def test_list_tensorboard_experiments_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -28925,11 +29303,15 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_experiment", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_experiment_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_experiment", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardExperimentRequest.pb( tensorboard_service.DeleteTensorboardExperimentRequest() ) @@ -28953,6 +29335,7 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_experiment( request, @@ -28964,6 +29347,7 @@ async def test_delete_tensorboard_experiment_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29152,10 +29536,14 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardRunRequest.pb( tensorboard_service.CreateTensorboardRunRequest() ) @@ -29181,6 +29569,7 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata await client.create_tensorboard_run( request, @@ -29192,6 +29581,7 @@ async def test_create_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29298,11 +29688,15 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_batch_create_tensorboard_runs", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_runs", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardRunsRequest.pb( tensorboard_service.BatchCreateTensorboardRunsRequest() ) @@ -29328,6 +29722,10 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.BatchCreateTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardRunsResponse(), + metadata, + ) await client.batch_create_tensorboard_runs( request, @@ -29339,6 +29737,7 @@ async def test_batch_create_tensorboard_runs_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29449,10 +29848,14 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardRunRequest.pb( tensorboard_service.GetTensorboardRunRequest() ) @@ -29478,6 +29881,7 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = tensorboard_run.TensorboardRun() + post_with_metadata.return_value = tensorboard_run.TensorboardRun(), metadata await client.get_tensorboard_run( request, @@ -29489,6 +29893,7 @@ async def test_get_tensorboard_run_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29681,10 +30086,14 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardRunRequest.pb( tensorboard_service.UpdateTensorboardRunRequest() ) @@ -29710,6 +30119,7 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = gca_tensorboard_run.TensorboardRun() + post_with_metadata.return_value = gca_tensorboard_run.TensorboardRun(), metadata await client.update_tensorboard_run( request, @@ -29721,6 +30131,7 @@ async def test_update_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29825,10 +30236,14 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_runs" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_runs_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_runs" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardRunsRequest.pb( tensorboard_service.ListTensorboardRunsRequest() ) @@ -29854,6 +30269,10 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardRunsResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardRunsResponse(), + metadata, + ) await client.list_tensorboard_runs( request, @@ -29865,6 +30284,7 @@ async def test_list_tensorboard_runs_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -29965,10 +30385,14 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor ), mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_run" ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_run_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_run" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardRunRequest.pb( tensorboard_service.DeleteTensorboardRunRequest() ) @@ -29992,6 +30416,7 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_run( request, @@ -30003,6 +30428,7 @@ async def test_delete_tensorboard_run_rest_asyncio_interceptors(null_interceptor pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30113,11 +30539,15 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_batch_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchCreateTensorboardTimeSeriesRequest.pb( tensorboard_service.BatchCreateTensorboardTimeSeriesRequest() ) @@ -30147,6 +30577,10 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( post.return_value = ( tensorboard_service.BatchCreateTensorboardTimeSeriesResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchCreateTensorboardTimeSeriesResponse(), + metadata, + ) await client.batch_create_tensorboard_time_series( request, @@ -30158,6 +30592,7 @@ async def test_batch_create_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30369,11 +30804,15 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_create_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_create_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_create_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.CreateTensorboardTimeSeriesRequest.pb( tensorboard_service.CreateTensorboardTimeSeriesRequest() ) @@ -30399,6 +30838,10 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.create_tensorboard_time_series( request, @@ -30410,6 +30853,7 @@ async def test_create_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30530,11 +30974,15 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc transports.AsyncTensorboardServiceRestInterceptor, "post_get_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_get_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_get_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.GetTensorboardTimeSeriesRequest.pb( tensorboard_service.GetTensorboardTimeSeriesRequest() ) @@ -30560,6 +31008,10 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc ] pre.return_value = request, metadata post.return_value = tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.get_tensorboard_time_series( request, @@ -30571,6 +31023,7 @@ async def test_get_tensorboard_time_series_rest_asyncio_interceptors(null_interc pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30786,11 +31239,15 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_update_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_update_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_update_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.UpdateTensorboardTimeSeriesRequest.pb( tensorboard_service.UpdateTensorboardTimeSeriesRequest() ) @@ -30816,6 +31273,10 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = gca_tensorboard_time_series.TensorboardTimeSeries() + post_with_metadata.return_value = ( + gca_tensorboard_time_series.TensorboardTimeSeries(), + metadata, + ) await client.update_tensorboard_time_series( request, @@ -30827,6 +31288,7 @@ async def test_update_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -30934,11 +31396,15 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter transports.AsyncTensorboardServiceRestInterceptor, "post_list_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_list_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_list_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ListTensorboardTimeSeriesRequest.pb( tensorboard_service.ListTensorboardTimeSeriesRequest() ) @@ -30964,6 +31430,10 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter ] pre.return_value = request, metadata post.return_value = tensorboard_service.ListTensorboardTimeSeriesResponse() + post_with_metadata.return_value = ( + tensorboard_service.ListTensorboardTimeSeriesResponse(), + metadata, + ) await client.list_tensorboard_time_series( request, @@ -30975,6 +31445,7 @@ async def test_list_tensorboard_time_series_rest_asyncio_interceptors(null_inter pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31078,11 +31549,15 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_delete_tensorboard_time_series", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_delete_tensorboard_time_series_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_delete_tensorboard_time_series", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.DeleteTensorboardTimeSeriesRequest.pb( tensorboard_service.DeleteTensorboardTimeSeriesRequest() ) @@ -31106,6 +31581,7 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_tensorboard_time_series( request, @@ -31117,6 +31593,7 @@ async def test_delete_tensorboard_time_series_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31229,11 +31706,15 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors transports.AsyncTensorboardServiceRestInterceptor, "post_batch_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_batch_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_batch_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.BatchReadTensorboardTimeSeriesDataRequest() ) @@ -31263,6 +31744,10 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors post.return_value = ( tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.BatchReadTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.batch_read_tensorboard_time_series_data( request, @@ -31274,6 +31759,7 @@ async def test_batch_read_tensorboard_time_series_data_rest_asyncio_interceptors pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31384,11 +31870,15 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ReadTensorboardTimeSeriesDataRequest() ) @@ -31416,6 +31906,10 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardTimeSeriesDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.read_tensorboard_time_series_data( request, @@ -31427,6 +31921,7 @@ async def test_read_tensorboard_time_series_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31533,11 +32028,15 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_read_tensorboard_blob_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_read_tensorboard_blob_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_read_tensorboard_blob_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ReadTensorboardBlobDataRequest.pb( tensorboard_service.ReadTensorboardBlobDataRequest() ) @@ -31563,6 +32062,10 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_service.ReadTensorboardBlobDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.ReadTensorboardBlobDataResponse(), + metadata, + ) await client.read_tensorboard_blob_data( request, @@ -31574,6 +32077,7 @@ async def test_read_tensorboard_blob_data_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31684,11 +32188,15 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_write_tensorboard_experiment_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_write_tensorboard_experiment_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_write_tensorboard_experiment_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardExperimentDataRequest.pb( tensorboard_service.WriteTensorboardExperimentDataRequest() ) @@ -31716,6 +32224,10 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardExperimentDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardExperimentDataResponse(), + metadata, + ) await client.write_tensorboard_experiment_data( request, @@ -31727,6 +32239,7 @@ async def test_write_tensorboard_experiment_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31831,11 +32344,15 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce transports.AsyncTensorboardServiceRestInterceptor, "post_write_tensorboard_run_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_write_tensorboard_run_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_write_tensorboard_run_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.WriteTensorboardRunDataRequest.pb( tensorboard_service.WriteTensorboardRunDataRequest() ) @@ -31861,6 +32378,10 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce ] pre.return_value = request, metadata post.return_value = tensorboard_service.WriteTensorboardRunDataResponse() + post_with_metadata.return_value = ( + tensorboard_service.WriteTensorboardRunDataResponse(), + metadata, + ) await client.write_tensorboard_run_data( request, @@ -31872,6 +32393,7 @@ async def test_write_tensorboard_run_data_rest_asyncio_interceptors(null_interce pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -31983,11 +32505,15 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( transports.AsyncTensorboardServiceRestInterceptor, "post_export_tensorboard_time_series_data", ) as post, mock.patch.object( + transports.AsyncTensorboardServiceRestInterceptor, + "post_export_tensorboard_time_series_data_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncTensorboardServiceRestInterceptor, "pre_export_tensorboard_time_series_data", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = tensorboard_service.ExportTensorboardTimeSeriesDataRequest.pb( tensorboard_service.ExportTensorboardTimeSeriesDataRequest() ) @@ -32017,6 +32543,10 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( post.return_value = ( tensorboard_service.ExportTensorboardTimeSeriesDataResponse() ) + post_with_metadata.return_value = ( + tensorboard_service.ExportTensorboardTimeSeriesDataResponse(), + metadata, + ) await client.export_tensorboard_time_series_data( request, @@ -32028,6 +32558,7 @@ async def test_export_tensorboard_time_series_data_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_data_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_data_service.py index d4e40803c2..649c828914 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_data_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_data_service.py @@ -91,6 +91,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -363,6 +371,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VertexRagDataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VertexRagDataServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -1868,6 +1919,7 @@ def test_get_rag_corpus(request_type, transport: str = "grpc"): name="name_value", display_name="display_name_value", description="description_value", + rag_files_count=1588, ) response = client.get_rag_corpus(request) @@ -1882,6 +1934,7 @@ def test_get_rag_corpus(request_type, transport: str = "grpc"): assert response.name == "name_value" assert response.display_name == "display_name_value" assert response.description == "description_value" + assert response.rag_files_count == 1588 def test_get_rag_corpus_non_empty_request_with_auto_populated_field(): @@ -2011,6 +2064,7 @@ async def test_get_rag_corpus_async( name="name_value", display_name="display_name_value", description="description_value", + rag_files_count=1588, ) ) response = await client.get_rag_corpus(request) @@ -2026,6 +2080,7 @@ async def test_get_rag_corpus_async( assert response.name == "name_value" assert response.display_name == "display_name_value" assert response.description == "description_value" + assert response.rag_files_count == 1588 @pytest.mark.asyncio @@ -7394,6 +7449,7 @@ async def test_get_rag_corpus_empty_call_grpc_asyncio(): name="name_value", display_name="display_name_value", description="description_value", + rag_files_count=1588, ) ) await client.get_rag_corpus(request=None) @@ -7682,6 +7738,7 @@ def test_create_rag_corpus_rest_call_success(request_type): "corpus_status": {"state": 1, "error_status": "error_status_value"}, "vector_db_config": {}, "vertex_ai_search_config": {"serving_config": "serving_config_value"}, + "rag_files_count": 1588, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -7791,10 +7848,14 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_create_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_create_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_create_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.CreateRagCorpusRequest.pb( vertex_rag_data_service.CreateRagCorpusRequest() ) @@ -7818,6 +7879,7 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.create_rag_corpus( request, @@ -7829,6 +7891,7 @@ def test_create_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_update_rag_corpus_rest_bad_request( @@ -7917,6 +7980,7 @@ def test_update_rag_corpus_rest_call_success(request_type): "corpus_status": {"state": 1, "error_status": "error_status_value"}, "vector_db_config": {}, "vertex_ai_search_config": {"serving_config": "serving_config_value"}, + "rag_files_count": 1588, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -8026,10 +8090,14 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_update_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_update_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_update_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UpdateRagCorpusRequest.pb( vertex_rag_data_service.UpdateRagCorpusRequest() ) @@ -8053,6 +8121,7 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.update_rag_corpus( request, @@ -8064,6 +8133,7 @@ def test_update_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_rag_corpus_rest_bad_request( @@ -8114,6 +8184,7 @@ def test_get_rag_corpus_rest_call_success(request_type): name="name_value", display_name="display_name_value", description="description_value", + rag_files_count=1588, ) # Wrap the value into a proper Response obj @@ -8133,6 +8204,7 @@ def test_get_rag_corpus_rest_call_success(request_type): assert response.name == "name_value" assert response.display_name == "display_name_value" assert response.description == "description_value" + assert response.rag_files_count == 1588 @pytest.mark.parametrize("null_interceptor", [True, False]) @@ -8152,10 +8224,14 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_get_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_get_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_get_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagCorpusRequest.pb( vertex_rag_data_service.GetRagCorpusRequest() ) @@ -8179,6 +8255,7 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagCorpus() + post_with_metadata.return_value = vertex_rag_data.RagCorpus(), metadata client.get_rag_corpus( request, @@ -8190,6 +8267,7 @@ def test_get_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_rag_corpora_rest_bad_request( @@ -8274,10 +8352,14 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_list_rag_corpora" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_list_rag_corpora_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_list_rag_corpora" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagCorporaRequest.pb( vertex_rag_data_service.ListRagCorporaRequest() ) @@ -8303,6 +8385,10 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagCorporaResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagCorporaResponse(), + metadata, + ) client.list_rag_corpora( request, @@ -8314,6 +8400,7 @@ def test_list_rag_corpora_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_rag_corpus_rest_bad_request( @@ -8394,10 +8481,14 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_delete_rag_corpus" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_delete_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_delete_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagCorpusRequest.pb( vertex_rag_data_service.DeleteRagCorpusRequest() ) @@ -8421,6 +8512,7 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_rag_corpus( request, @@ -8432,6 +8524,7 @@ def test_delete_rag_corpus_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_upload_rag_file_rest_bad_request( @@ -8513,10 +8606,14 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_upload_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_upload_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_upload_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UploadRagFileRequest.pb( vertex_rag_data_service.UploadRagFileRequest() ) @@ -8542,6 +8639,10 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.UploadRagFileResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.UploadRagFileResponse(), + metadata, + ) client.upload_rag_file( request, @@ -8553,6 +8654,7 @@ def test_upload_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_import_rag_files_rest_bad_request( @@ -8633,10 +8735,14 @@ def test_import_rag_files_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_import_rag_files" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_import_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_import_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ImportRagFilesRequest.pb( vertex_rag_data_service.ImportRagFilesRequest() ) @@ -8660,6 +8766,7 @@ def test_import_rag_files_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.import_rag_files( request, @@ -8671,6 +8778,7 @@ def test_import_rag_files_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_rag_file_rest_bad_request( @@ -8769,10 +8877,14 @@ def test_get_rag_file_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_get_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_get_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_get_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagFileRequest.pb( vertex_rag_data_service.GetRagFileRequest() ) @@ -8796,6 +8908,7 @@ def test_get_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagFile() + post_with_metadata.return_value = vertex_rag_data.RagFile(), metadata client.get_rag_file( request, @@ -8807,6 +8920,7 @@ def test_get_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_rag_files_rest_bad_request( @@ -8891,10 +9005,14 @@ def test_list_rag_files_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_list_rag_files" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_list_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_list_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagFilesRequest.pb( vertex_rag_data_service.ListRagFilesRequest() ) @@ -8920,6 +9038,10 @@ def test_list_rag_files_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagFilesResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagFilesResponse(), + metadata, + ) client.list_rag_files( request, @@ -8931,6 +9053,7 @@ def test_list_rag_files_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_rag_file_rest_bad_request( @@ -9015,10 +9138,14 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "post_delete_rag_file" ) as post, mock.patch.object( + transports.VertexRagDataServiceRestInterceptor, + "post_delete_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagDataServiceRestInterceptor, "pre_delete_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagFileRequest.pb( vertex_rag_data_service.DeleteRagFileRequest() ) @@ -9042,6 +9169,7 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.delete_rag_file( request, @@ -9053,6 +9181,7 @@ def test_delete_rag_file_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -10012,6 +10141,7 @@ async def test_create_rag_corpus_rest_asyncio_call_success(request_type): "corpus_status": {"state": 1, "error_status": "error_status_value"}, "vector_db_config": {}, "vertex_ai_search_config": {"serving_config": "serving_config_value"}, + "rag_files_count": 1588, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -10128,10 +10258,14 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_create_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_create_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_create_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.CreateRagCorpusRequest.pb( vertex_rag_data_service.CreateRagCorpusRequest() ) @@ -10155,6 +10289,7 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.create_rag_corpus( request, @@ -10166,6 +10301,7 @@ async def test_create_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10263,6 +10399,7 @@ async def test_update_rag_corpus_rest_asyncio_call_success(request_type): "corpus_status": {"state": 1, "error_status": "error_status_value"}, "vector_db_config": {}, "vertex_ai_search_config": {"serving_config": "serving_config_value"}, + "rag_files_count": 1588, } # The version of a generated dependency at test runtime may differ from the version used during generation. # Delete any fields which are not present in the current runtime dependency @@ -10379,10 +10516,14 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_update_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_update_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_update_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UpdateRagCorpusRequest.pb( vertex_rag_data_service.UpdateRagCorpusRequest() ) @@ -10406,6 +10547,7 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.update_rag_corpus( request, @@ -10417,6 +10559,7 @@ async def test_update_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10476,6 +10619,7 @@ async def test_get_rag_corpus_rest_asyncio_call_success(request_type): name="name_value", display_name="display_name_value", description="description_value", + rag_files_count=1588, ) # Wrap the value into a proper Response obj @@ -10497,6 +10641,7 @@ async def test_get_rag_corpus_rest_asyncio_call_success(request_type): assert response.name == "name_value" assert response.display_name == "display_name_value" assert response.description == "description_value" + assert response.rag_files_count == 1588 @pytest.mark.asyncio @@ -10521,10 +10666,14 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_get_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_get_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_get_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagCorpusRequest.pb( vertex_rag_data_service.GetRagCorpusRequest() ) @@ -10548,6 +10697,7 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagCorpus() + post_with_metadata.return_value = vertex_rag_data.RagCorpus(), metadata await client.get_rag_corpus( request, @@ -10559,6 +10709,7 @@ async def test_get_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10659,10 +10810,14 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_list_rag_corpora" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_list_rag_corpora_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_list_rag_corpora" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagCorporaRequest.pb( vertex_rag_data_service.ListRagCorporaRequest() ) @@ -10688,6 +10843,10 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagCorporaResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagCorporaResponse(), + metadata, + ) await client.list_rag_corpora( request, @@ -10699,6 +10858,7 @@ async def test_list_rag_corpora_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10795,10 +10955,14 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_delete_rag_corpus" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_delete_rag_corpus_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_delete_rag_corpus" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagCorpusRequest.pb( vertex_rag_data_service.DeleteRagCorpusRequest() ) @@ -10822,6 +10986,7 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_rag_corpus( request, @@ -10833,6 +10998,7 @@ async def test_delete_rag_corpus_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -10930,10 +11096,14 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_upload_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_upload_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_upload_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.UploadRagFileRequest.pb( vertex_rag_data_service.UploadRagFileRequest() ) @@ -10959,6 +11129,10 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.UploadRagFileResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.UploadRagFileResponse(), + metadata, + ) await client.upload_rag_file( request, @@ -10970,6 +11144,7 @@ async def test_upload_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11066,10 +11241,14 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_import_rag_files" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_import_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_import_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ImportRagFilesRequest.pb( vertex_rag_data_service.ImportRagFilesRequest() ) @@ -11093,6 +11272,7 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.import_rag_files( request, @@ -11104,6 +11284,7 @@ async def test_import_rag_files_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11218,10 +11399,14 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_get_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_get_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_get_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.GetRagFileRequest.pb( vertex_rag_data_service.GetRagFileRequest() ) @@ -11245,6 +11430,7 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data.RagFile() + post_with_metadata.return_value = vertex_rag_data.RagFile(), metadata await client.get_rag_file( request, @@ -11256,6 +11442,7 @@ async def test_get_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11356,10 +11543,14 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_list_rag_files" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_list_rag_files_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_list_rag_files" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.ListRagFilesRequest.pb( vertex_rag_data_service.ListRagFilesRequest() ) @@ -11385,6 +11576,10 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_data_service.ListRagFilesResponse() + post_with_metadata.return_value = ( + vertex_rag_data_service.ListRagFilesResponse(), + metadata, + ) await client.list_rag_files( request, @@ -11396,6 +11591,7 @@ async def test_list_rag_files_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -11496,10 +11692,14 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "post_delete_rag_file" ) as post, mock.patch.object( + transports.AsyncVertexRagDataServiceRestInterceptor, + "post_delete_rag_file_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagDataServiceRestInterceptor, "pre_delete_rag_file" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_data_service.DeleteRagFileRequest.pb( vertex_rag_data_service.DeleteRagFileRequest() ) @@ -11523,6 +11723,7 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.delete_rag_file( request, @@ -11534,6 +11735,7 @@ async def test_delete_rag_file_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_service.py index 57e83323f5..73bca4cd81 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_vertex_rag_service.py @@ -84,6 +84,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -342,6 +350,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VertexRagServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VertexRagServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -3073,10 +3124,14 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_retrieve_contexts" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, + "post_retrieve_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_retrieve_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.RetrieveContextsRequest.pb( vertex_rag_service.RetrieveContextsRequest() ) @@ -3102,6 +3157,10 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.RetrieveContextsResponse() + post_with_metadata.return_value = ( + vertex_rag_service.RetrieveContextsResponse(), + metadata, + ) client.retrieve_contexts( request, @@ -3113,6 +3172,7 @@ def test_retrieve_contexts_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_augment_prompt_rest_bad_request( @@ -3194,10 +3254,13 @@ def test_augment_prompt_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_augment_prompt" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, "post_augment_prompt_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_augment_prompt" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.AugmentPromptRequest.pb( vertex_rag_service.AugmentPromptRequest() ) @@ -3223,6 +3286,10 @@ def test_augment_prompt_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.AugmentPromptResponse() + post_with_metadata.return_value = ( + vertex_rag_service.AugmentPromptResponse(), + metadata, + ) client.augment_prompt( request, @@ -3234,6 +3301,7 @@ def test_augment_prompt_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_corroborate_content_rest_bad_request( @@ -3318,10 +3386,14 @@ def test_corroborate_content_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VertexRagServiceRestInterceptor, "post_corroborate_content" ) as post, mock.patch.object( + transports.VertexRagServiceRestInterceptor, + "post_corroborate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VertexRagServiceRestInterceptor, "pre_corroborate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.CorroborateContentRequest.pb( vertex_rag_service.CorroborateContentRequest() ) @@ -3347,6 +3419,10 @@ def test_corroborate_content_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.CorroborateContentResponse() + post_with_metadata.return_value = ( + vertex_rag_service.CorroborateContentResponse(), + metadata, + ) client.corroborate_content( request, @@ -3358,6 +3434,7 @@ def test_corroborate_content_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -4162,10 +4239,14 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_retrieve_contexts" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_retrieve_contexts_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_retrieve_contexts" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.RetrieveContextsRequest.pb( vertex_rag_service.RetrieveContextsRequest() ) @@ -4191,6 +4272,10 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.RetrieveContextsResponse() + post_with_metadata.return_value = ( + vertex_rag_service.RetrieveContextsResponse(), + metadata, + ) await client.retrieve_contexts( request, @@ -4202,6 +4287,7 @@ async def test_retrieve_contexts_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4299,10 +4385,14 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_augment_prompt" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_augment_prompt_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_augment_prompt" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.AugmentPromptRequest.pb( vertex_rag_service.AugmentPromptRequest() ) @@ -4328,6 +4418,10 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.AugmentPromptResponse() + post_with_metadata.return_value = ( + vertex_rag_service.AugmentPromptResponse(), + metadata, + ) await client.augment_prompt( request, @@ -4339,6 +4433,7 @@ async def test_augment_prompt_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -4439,10 +4534,14 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "post_corroborate_content" ) as post, mock.patch.object( + transports.AsyncVertexRagServiceRestInterceptor, + "post_corroborate_content_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVertexRagServiceRestInterceptor, "pre_corroborate_content" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vertex_rag_service.CorroborateContentRequest.pb( vertex_rag_service.CorroborateContentRequest() ) @@ -4468,6 +4567,10 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vertex_rag_service.CorroborateContentResponse() + post_with_metadata.return_value = ( + vertex_rag_service.CorroborateContentResponse(), + metadata, + ) await client.corroborate_content( request, @@ -4479,6 +4582,7 @@ async def test_corroborate_content_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/gapic/aiplatform_v1beta1/test_vizier_service.py b/tests/unit/gapic/aiplatform_v1beta1/test_vizier_service.py index 71259bd737..586aeb52f9 100644 --- a/tests/unit/gapic/aiplatform_v1beta1/test_vizier_service.py +++ b/tests/unit/gapic/aiplatform_v1beta1/test_vizier_service.py @@ -88,6 +88,14 @@ import google.auth +CRED_INFO_JSON = { + "credential_source": "/path/to/file", + "credential_type": "service account credentials", + "principal": "service-account@example.com", +} +CRED_INFO_STRING = json.dumps(CRED_INFO_JSON) + + async def mock_async_gen(data, chunk_size=1): for i in range(0, len(data)): # pragma: NO COVER chunk = data[i : i + chunk_size] @@ -341,6 +349,49 @@ def test__get_universe_domain(): assert str(excinfo.value) == "Universe Domain cannot be an empty string." +@pytest.mark.parametrize( + "error_code,cred_info_json,show_cred_info", + [ + (401, CRED_INFO_JSON, True), + (403, CRED_INFO_JSON, True), + (404, CRED_INFO_JSON, True), + (500, CRED_INFO_JSON, False), + (401, None, False), + (403, None, False), + (404, None, False), + (500, None, False), + ], +) +def test__add_cred_info_for_auth_errors(error_code, cred_info_json, show_cred_info): + cred = mock.Mock(["get_cred_info"]) + cred.get_cred_info = mock.Mock(return_value=cred_info_json) + client = VizierServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=["foo"]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + if show_cred_info: + assert error.details == ["foo", CRED_INFO_STRING] + else: + assert error.details == ["foo"] + + +@pytest.mark.parametrize("error_code", [401, 403, 404, 500]) +def test__add_cred_info_for_auth_errors_no_get_cred_info(error_code): + cred = mock.Mock([]) + assert not hasattr(cred, "get_cred_info") + client = VizierServiceClient(credentials=cred) + client._transport._credentials = cred + + error = core_exceptions.GoogleAPICallError("message", details=[]) + error.code = error_code + + client._add_cred_info_for_auth_errors(error) + assert error.details == [] + + @pytest.mark.parametrize( "client_class,transport_name", [ @@ -9807,10 +9858,13 @@ def test_create_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_create_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_create_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_create_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateStudyRequest.pb( vizier_service.CreateStudyRequest() ) @@ -9834,6 +9888,7 @@ def test_create_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_study.Study() + post_with_metadata.return_value = gca_study.Study(), metadata client.create_study( request, @@ -9845,6 +9900,7 @@ def test_create_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_study_rest_bad_request(request_type=vizier_service.GetStudyRequest): @@ -9933,10 +9989,13 @@ def test_get_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_get_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_get_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_get_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetStudyRequest.pb(vizier_service.GetStudyRequest()) transcode.return_value = { "method": "post", @@ -9958,6 +10017,7 @@ def test_get_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata client.get_study( request, @@ -9969,6 +10029,7 @@ def test_get_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_studies_rest_bad_request(request_type=vizier_service.ListStudiesRequest): @@ -10051,10 +10112,13 @@ def test_list_studies_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_studies" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_list_studies_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_studies" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListStudiesRequest.pb( vizier_service.ListStudiesRequest() ) @@ -10080,6 +10144,7 @@ def test_list_studies_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListStudiesResponse() + post_with_metadata.return_value = vizier_service.ListStudiesResponse(), metadata client.list_studies( request, @@ -10091,6 +10156,7 @@ def test_list_studies_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_study_rest_bad_request(request_type=vizier_service.DeleteStudyRequest): @@ -10286,10 +10352,13 @@ def test_lookup_study_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_lookup_study" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_lookup_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_lookup_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.LookupStudyRequest.pb( vizier_service.LookupStudyRequest() ) @@ -10313,6 +10382,7 @@ def test_lookup_study_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata client.lookup_study( request, @@ -10324,6 +10394,7 @@ def test_lookup_study_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_suggest_trials_rest_bad_request( @@ -10404,10 +10475,13 @@ def test_suggest_trials_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VizierServiceRestInterceptor, "post_suggest_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_suggest_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_suggest_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.SuggestTrialsRequest.pb( vizier_service.SuggestTrialsRequest() ) @@ -10431,6 +10505,7 @@ def test_suggest_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.suggest_trials( request, @@ -10442,6 +10517,7 @@ def test_suggest_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_create_trial_rest_bad_request(request_type=vizier_service.CreateTrialRequest): @@ -10631,10 +10707,13 @@ def test_create_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_create_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_create_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_create_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateTrialRequest.pb( vizier_service.CreateTrialRequest() ) @@ -10658,6 +10737,7 @@ def test_create_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.create_trial( request, @@ -10669,6 +10749,7 @@ def test_create_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_trial_rest_bad_request(request_type=vizier_service.GetTrialRequest): @@ -10765,10 +10846,13 @@ def test_get_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_get_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_get_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_get_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetTrialRequest.pb(vizier_service.GetTrialRequest()) transcode.return_value = { "method": "post", @@ -10790,6 +10874,7 @@ def test_get_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.get_trial( request, @@ -10801,6 +10886,7 @@ def test_get_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_trials_rest_bad_request(request_type=vizier_service.ListTrialsRequest): @@ -10883,10 +10969,13 @@ def test_list_trials_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_list_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListTrialsRequest.pb( vizier_service.ListTrialsRequest() ) @@ -10912,6 +11001,7 @@ def test_list_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListTrialsResponse() + post_with_metadata.return_value = vizier_service.ListTrialsResponse(), metadata client.list_trials( request, @@ -10923,6 +11013,7 @@ def test_list_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_add_trial_measurement_rest_bad_request( @@ -11021,10 +11112,14 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_add_trial_measurement" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_add_trial_measurement_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_add_trial_measurement" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.AddTrialMeasurementRequest.pb( vizier_service.AddTrialMeasurementRequest() ) @@ -11048,6 +11143,7 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.add_trial_measurement( request, @@ -11059,6 +11155,7 @@ def test_add_trial_measurement_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_complete_trial_rest_bad_request( @@ -11157,10 +11254,13 @@ def test_complete_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_complete_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_complete_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_complete_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CompleteTrialRequest.pb( vizier_service.CompleteTrialRequest() ) @@ -11184,6 +11284,7 @@ def test_complete_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.complete_trial( request, @@ -11195,6 +11296,7 @@ def test_complete_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_delete_trial_rest_bad_request(request_type=vizier_service.DeleteTrialRequest): @@ -11390,10 +11492,14 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): ), mock.patch.object( transports.VizierServiceRestInterceptor, "post_check_trial_early_stopping_state" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_check_trial_early_stopping_state_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_check_trial_early_stopping_state" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CheckTrialEarlyStoppingStateRequest.pb( vizier_service.CheckTrialEarlyStoppingStateRequest() ) @@ -11417,6 +11523,7 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata client.check_trial_early_stopping_state( request, @@ -11428,6 +11535,7 @@ def test_check_trial_early_stopping_state_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_stop_trial_rest_bad_request(request_type=vizier_service.StopTrialRequest): @@ -11524,10 +11632,13 @@ def test_stop_trial_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_stop_trial" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, "post_stop_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_stop_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.StopTrialRequest.pb( vizier_service.StopTrialRequest() ) @@ -11551,6 +11662,7 @@ def test_stop_trial_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata client.stop_trial( request, @@ -11562,6 +11674,7 @@ def test_stop_trial_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_list_optimal_trials_rest_bad_request( @@ -11643,10 +11756,14 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.VizierServiceRestInterceptor, "post_list_optimal_trials" ) as post, mock.patch.object( + transports.VizierServiceRestInterceptor, + "post_list_optimal_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.VizierServiceRestInterceptor, "pre_list_optimal_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListOptimalTrialsRequest.pb( vizier_service.ListOptimalTrialsRequest() ) @@ -11672,6 +11789,10 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListOptimalTrialsResponse() + post_with_metadata.return_value = ( + vizier_service.ListOptimalTrialsResponse(), + metadata, + ) client.list_optimal_trials( request, @@ -11683,6 +11804,7 @@ def test_list_optimal_trials_rest_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() def test_get_location_rest_bad_request(request_type=locations_pb2.GetLocationRequest): @@ -12914,10 +13036,13 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_create_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_create_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_create_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateStudyRequest.pb( vizier_service.CreateStudyRequest() ) @@ -12941,6 +13066,7 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = gca_study.Study() + post_with_metadata.return_value = gca_study.Study(), metadata await client.create_study( request, @@ -12952,6 +13078,7 @@ async def test_create_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13058,10 +13185,13 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_get_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_get_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_get_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetStudyRequest.pb(vizier_service.GetStudyRequest()) transcode.return_value = { "method": "post", @@ -13083,6 +13213,7 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata await client.get_study( request, @@ -13094,6 +13225,7 @@ async def test_get_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13194,10 +13326,13 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_studies" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_list_studies_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_studies" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListStudiesRequest.pb( vizier_service.ListStudiesRequest() ) @@ -13223,6 +13358,7 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListStudiesResponse() + post_with_metadata.return_value = vizier_service.ListStudiesResponse(), metadata await client.list_studies( request, @@ -13234,6 +13370,7 @@ async def test_list_studies_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13465,10 +13602,13 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_lookup_study" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_lookup_study_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_lookup_study" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.LookupStudyRequest.pb( vizier_service.LookupStudyRequest() ) @@ -13492,6 +13632,7 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Study() + post_with_metadata.return_value = study.Study(), metadata await client.lookup_study( request, @@ -13503,6 +13644,7 @@ async def test_lookup_study_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13599,10 +13741,14 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): ), mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_suggest_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_suggest_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_suggest_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.SuggestTrialsRequest.pb( vizier_service.SuggestTrialsRequest() ) @@ -13626,6 +13772,7 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.suggest_trials( request, @@ -13637,6 +13784,7 @@ async def test_suggest_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13844,10 +13992,13 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_create_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_create_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_create_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CreateTrialRequest.pb( vizier_service.CreateTrialRequest() ) @@ -13871,6 +14022,7 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.create_trial( request, @@ -13882,6 +14034,7 @@ async def test_create_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -13996,10 +14149,13 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_get_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_get_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_get_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.GetTrialRequest.pb(vizier_service.GetTrialRequest()) transcode.return_value = { "method": "post", @@ -14021,6 +14177,7 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.get_trial( request, @@ -14032,6 +14189,7 @@ async def test_get_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14132,10 +14290,13 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_list_trials_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListTrialsRequest.pb( vizier_service.ListTrialsRequest() ) @@ -14161,6 +14322,7 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListTrialsResponse() + post_with_metadata.return_value = vizier_service.ListTrialsResponse(), metadata await client.list_trials( request, @@ -14172,6 +14334,7 @@ async def test_list_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14286,10 +14449,14 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_add_trial_measurement" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_add_trial_measurement_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_add_trial_measurement" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.AddTrialMeasurementRequest.pb( vizier_service.AddTrialMeasurementRequest() ) @@ -14313,6 +14480,7 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.add_trial_measurement( request, @@ -14324,6 +14492,7 @@ async def test_add_trial_measurement_rest_asyncio_interceptors(null_interceptor) pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14438,10 +14607,14 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_complete_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_complete_trial_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_complete_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CompleteTrialRequest.pb( vizier_service.CompleteTrialRequest() ) @@ -14465,6 +14638,7 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.complete_trial( request, @@ -14476,6 +14650,7 @@ async def test_complete_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14708,11 +14883,15 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( transports.AsyncVizierServiceRestInterceptor, "post_check_trial_early_stopping_state", ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_check_trial_early_stopping_state_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_check_trial_early_stopping_state", ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.CheckTrialEarlyStoppingStateRequest.pb( vizier_service.CheckTrialEarlyStoppingStateRequest() ) @@ -14736,6 +14915,7 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( ] pre.return_value = request, metadata post.return_value = operations_pb2.Operation() + post_with_metadata.return_value = operations_pb2.Operation(), metadata await client.check_trial_early_stopping_state( request, @@ -14747,6 +14927,7 @@ async def test_check_trial_early_stopping_state_rest_asyncio_interceptors( pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14861,10 +15042,13 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_stop_trial" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, "post_stop_trial_with_metadata" + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_stop_trial" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.StopTrialRequest.pb( vizier_service.StopTrialRequest() ) @@ -14888,6 +15072,7 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = study.Trial() + post_with_metadata.return_value = study.Trial(), metadata await client.stop_trial( request, @@ -14899,6 +15084,7 @@ async def test_stop_trial_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio @@ -14996,10 +15182,14 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): ) as transcode, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "post_list_optimal_trials" ) as post, mock.patch.object( + transports.AsyncVizierServiceRestInterceptor, + "post_list_optimal_trials_with_metadata", + ) as post_with_metadata, mock.patch.object( transports.AsyncVizierServiceRestInterceptor, "pre_list_optimal_trials" ) as pre: pre.assert_not_called() post.assert_not_called() + post_with_metadata.assert_not_called() pb_message = vizier_service.ListOptimalTrialsRequest.pb( vizier_service.ListOptimalTrialsRequest() ) @@ -15025,6 +15215,10 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): ] pre.return_value = request, metadata post.return_value = vizier_service.ListOptimalTrialsResponse() + post_with_metadata.return_value = ( + vizier_service.ListOptimalTrialsResponse(), + metadata, + ) await client.list_optimal_trials( request, @@ -15036,6 +15230,7 @@ async def test_list_optimal_trials_rest_asyncio_interceptors(null_interceptor): pre.assert_called_once() post.assert_called_once() + post_with_metadata.assert_called_once() @pytest.mark.asyncio diff --git a/tests/unit/vertexai/test_evaluation.py b/tests/unit/vertexai/test_evaluation.py index 376c4e7944..364f81df34 100644 --- a/tests/unit/vertexai/test_evaluation.py +++ b/tests/unit/vertexai/test_evaluation.py @@ -25,6 +25,7 @@ from google.cloud import aiplatform import vertexai from google.cloud.aiplatform import initializer +from google.cloud.aiplatform import utils as aiplatform_utils from google.cloud.aiplatform.metadata import metadata from google.cloud.aiplatform_v1.services import ( evaluation_service as gapic_evaluation_services, @@ -2176,3 +2177,18 @@ def test_upload_results(self, mock_storage_blob_from_string): }, mock.ANY, ) + + def test_upload_results_with_default_file_name(self, mock_storage_blob_from_string): + with mock.patch.object( + aiplatform_utils, "timestamped_unique_name" + ) as mock_timestamped_unique_name: + mock_timestamped_unique_name.return_value = "2025-02-10-12-00-00-12345" + evaluation.utils.upload_evaluation_results( + MOCK_EVAL_RESULT, + _TEST_BUCKET, + ) + + mock_storage_blob_from_string.assert_any_call( + uri="gs://test-bucket/eval_results_2025-02-10-12-00-00-12345/eval_results_2025-02-10-12-00-00-12345.csv", + client=mock.ANY, + ) diff --git a/vertexai/evaluation/utils.py b/vertexai/evaluation/utils.py index ff5df99b3b..2c2c696340 100644 --- a/vertexai/evaluation/utils.py +++ b/vertexai/evaluation/utils.py @@ -336,7 +336,7 @@ def _upload_file_to_gcs(upload_gcs_path: str, filename: str) -> None: def upload_evaluation_results( eval_result: eval_base.EvalResult, destination_uri_prefix: str, - file_name: str, + file_name: Optional[str] = None, candidate_model_name: Optional[str] = None, baseline_model_name: Optional[str] = None, dataset_uri: Optional[str] = None, @@ -347,7 +347,7 @@ def upload_evaluation_results( Args: eval_result: Eval results to upload. destination_uri_prefix: GCS folder to store the data. - file_name: File name to store the metrics table. + file_name: Optional. File name to store the metrics table. candidate_model_name: Optional. Candidate model name. baseline_model_name: Optional. Baseline model name. dataset_uri: Optional. URI pointing to the dataset. @@ -359,6 +359,8 @@ def upload_evaluation_results( if eval_result.metrics_table is None: return if destination_uri_prefix.startswith(_GCS_PREFIX): + if not file_name: + file_name = f"eval_results_{utils.timestamped_unique_name()}.csv" base_name, extension = os.path.splitext(file_name) file_type = extension.lower()[1:] output_folder = destination_uri_prefix + "/" + base_name diff --git a/vertexai/vision_models/_vision_models.py b/vertexai/vision_models/_vision_models.py index 1454b2b711..c252fc3e6f 100644 --- a/vertexai/vision_models/_vision_models.py +++ b/vertexai/vision_models/_vision_models.py @@ -1938,6 +1938,7 @@ def segment_image( max_predictions: Optional[int] = None, confidence_threshold: Optional[float] = 0.1, mask_dilation: Optional[float] = None, + binary_color_threshold: Optional[float] = None, ) -> ImageSegmentationResponse: """Segments an image. @@ -1967,6 +1968,10 @@ def segment_image( mask_dilation: A value to dilate the masks by. The value must be in the range of 0.0 (no dilation) and 1.0 (the whole image will be masked). The default is 0.0. + binary_color_threshold: The threshold to convert the grayscale soft + mask to a binary color black and white mask. The value must be + in the range of 0 and 255, or -1 to disable the thresholding. + The default is 96. Returns: An `ImageSegmentationResponse` object with the generated masks, @@ -2000,6 +2005,8 @@ def segment_image( parameters["confidenceThreshold"] = confidence_threshold if mask_dilation: parameters["maskDilation"] = mask_dilation + if binary_color_threshold: + parameters["binaryColorThreshold"] = binary_color_threshold response = self._endpoint.predict( instances=[instance],