diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service/__init__.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service/__init__.py index b79f7274cddc..962fbf440f12 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service/__init__.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service/__init__.py @@ -31,6 +31,8 @@ ImageVersionsClient, ) from google.cloud.orchestration.airflow.service_v1.types.environments import ( + AirflowMetadataRetentionPolicyConfig, + CheckUpgradeRequest, CheckUpgradeResponse, CloudDataLineageIntegration, CreateEnvironmentRequest, @@ -104,6 +106,8 @@ "EnvironmentsAsyncClient", "ImageVersionsClient", "ImageVersionsAsyncClient", + "AirflowMetadataRetentionPolicyConfig", + "CheckUpgradeRequest", "CheckUpgradeResponse", "CloudDataLineageIntegration", "CreateEnvironmentRequest", diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/__init__.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/__init__.py index a30d9c61ae59..2cab45afb1ee 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/__init__.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/__init__.py @@ -23,6 +23,8 @@ from .services.environments import EnvironmentsAsyncClient, EnvironmentsClient from .services.image_versions import ImageVersionsAsyncClient, ImageVersionsClient from .types.environments import ( + AirflowMetadataRetentionPolicyConfig, + CheckUpgradeRequest, CheckUpgradeResponse, CloudDataLineageIntegration, CreateEnvironmentRequest, @@ -92,6 +94,8 @@ __all__ = ( "EnvironmentsAsyncClient", "ImageVersionsAsyncClient", + "AirflowMetadataRetentionPolicyConfig", + "CheckUpgradeRequest", "CheckUpgradeResponse", "CloudDataLineageIntegration", "CreateEnvironmentRequest", diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/gapic_metadata.json b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/gapic_metadata.json index 3511b1f8d28f..aca0f03a3f15 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/gapic_metadata.json +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/gapic_metadata.json @@ -10,6 +10,11 @@ "grpc": { "libraryClient": "EnvironmentsClient", "rpcs": { + "CheckUpgrade": { + "methods": [ + "check_upgrade" + ] + }, "CreateEnvironment": { "methods": [ "create_environment" @@ -130,6 +135,11 @@ "grpc-async": { "libraryClient": "EnvironmentsAsyncClient", "rpcs": { + "CheckUpgrade": { + "methods": [ + "check_upgrade" + ] + }, "CreateEnvironment": { "methods": [ "create_environment" @@ -250,6 +260,11 @@ "rest": { "libraryClient": "EnvironmentsClient", "rpcs": { + "CheckUpgrade": { + "methods": [ + "check_upgrade" + ] + }, "CreateEnvironment": { "methods": [ "create_environment" diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/async_client.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/async_client.py index 3406e2214500..849eed1fbf2d 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/async_client.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/async_client.py @@ -1448,6 +1448,109 @@ async def sample_list_workloads(): # Done; return the response. return response + async def check_upgrade( + self, + request: Optional[Union[environments.CheckUpgradeRequest, dict]] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operation_async.AsyncOperation: + r"""Check if an upgrade operation on the environment will + succeed. + In case of problems detailed info can be found in the + returned Operation. + + .. 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.orchestration.airflow import service_v1 + + async def sample_check_upgrade(): + # Create a client + client = service_v1.EnvironmentsAsyncClient() + + # Initialize request argument(s) + request = service_v1.CheckUpgradeRequest( + environment="environment_value", + ) + + # Make the request + operation = client.check_upgrade(request=request) + + print("Waiting for operation to complete...") + + response = (await operation).result() + + # Handle the response + print(response) + + Args: + request (Optional[Union[google.cloud.orchestration.airflow.service_v1.types.CheckUpgradeRequest, dict]]): + The request object. Request to check whether image + upgrade will succeed. + 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, str]]): Strings which should be + sent along with the request as metadata. + + 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.orchestration.airflow.service_v1.types.CheckUpgradeResponse` Message containing information about the result of an upgrade check + operation. + + """ + # 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, environments.CheckUpgradeRequest): + request = environments.CheckUpgradeRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._client._transport._wrapped_methods[ + self._client._transport.check_upgrade + ] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("environment", request.environment),) + ), + ) + + # 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, + environments.CheckUpgradeResponse, + metadata_type=operations.OperationMetadata, + ) + + # Done; return the response. + return response + async def create_user_workloads_secret( self, request: Optional[ diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/client.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/client.py index a00e4bf1e986..6a8635aa5380 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/client.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/client.py @@ -1884,6 +1884,107 @@ def sample_list_workloads(): # Done; return the response. return response + def check_upgrade( + self, + request: Optional[Union[environments.CheckUpgradeRequest, dict]] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operation.Operation: + r"""Check if an upgrade operation on the environment will + succeed. + In case of problems detailed info can be found in the + returned Operation. + + .. 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.orchestration.airflow import service_v1 + + def sample_check_upgrade(): + # Create a client + client = service_v1.EnvironmentsClient() + + # Initialize request argument(s) + request = service_v1.CheckUpgradeRequest( + environment="environment_value", + ) + + # Make the request + operation = client.check_upgrade(request=request) + + print("Waiting for operation to complete...") + + response = operation.result() + + # Handle the response + print(response) + + Args: + request (Union[google.cloud.orchestration.airflow.service_v1.types.CheckUpgradeRequest, dict]): + The request object. Request to check whether image + upgrade will succeed. + 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, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.api_core.operation.Operation: + An object representing a long-running operation. + + The result type for the operation will be :class:`google.cloud.orchestration.airflow.service_v1.types.CheckUpgradeResponse` Message containing information about the result of an upgrade check + operation. + + """ + # 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, environments.CheckUpgradeRequest): + request = environments.CheckUpgradeRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.check_upgrade] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("environment", request.environment),) + ), + ) + + # 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, + environments.CheckUpgradeResponse, + metadata_type=operations.OperationMetadata, + ) + + # Done; return the response. + return response + def create_user_workloads_secret( self, request: Optional[ diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/base.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/base.py index 735acee260d9..05b1c3627955 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/base.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/base.py @@ -177,6 +177,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.check_upgrade: gapic_v1.method.wrap_method( + self.check_upgrade, + default_timeout=None, + client_info=client_info, + ), self.create_user_workloads_secret: gapic_v1.method.wrap_method( self.create_user_workloads_secret, default_timeout=None, @@ -359,6 +364,15 @@ def list_workloads( ]: raise NotImplementedError() + @property + def check_upgrade( + self, + ) -> Callable[ + [environments.CheckUpgradeRequest], + Union[operations_pb2.Operation, Awaitable[operations_pb2.Operation]], + ]: + raise NotImplementedError() + @property def create_user_workloads_secret( self, diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc.py index 78c4216a1d96..1d424d9d8162 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc.py @@ -502,6 +502,35 @@ def list_workloads( ) return self._stubs["list_workloads"] + @property + def check_upgrade( + self, + ) -> Callable[[environments.CheckUpgradeRequest], operations_pb2.Operation]: + r"""Return a callable for the check upgrade method over gRPC. + + Check if an upgrade operation on the environment will + succeed. + In case of problems detailed info can be found in the + returned Operation. + + Returns: + Callable[[~.CheckUpgradeRequest], + ~.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 "check_upgrade" not in self._stubs: + self._stubs["check_upgrade"] = self.grpc_channel.unary_unary( + "/google.cloud.orchestration.airflow.service.v1.Environments/CheckUpgrade", + request_serializer=environments.CheckUpgradeRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["check_upgrade"] + @property def create_user_workloads_secret( self, diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc_asyncio.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc_asyncio.py index cbe9ca640a9c..0c408a50e4cf 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc_asyncio.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/grpc_asyncio.py @@ -518,6 +518,37 @@ def list_workloads( ) return self._stubs["list_workloads"] + @property + def check_upgrade( + self, + ) -> Callable[ + [environments.CheckUpgradeRequest], Awaitable[operations_pb2.Operation] + ]: + r"""Return a callable for the check upgrade method over gRPC. + + Check if an upgrade operation on the environment will + succeed. + In case of problems detailed info can be found in the + returned Operation. + + Returns: + Callable[[~.CheckUpgradeRequest], + 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 "check_upgrade" not in self._stubs: + self._stubs["check_upgrade"] = self.grpc_channel.unary_unary( + "/google.cloud.orchestration.airflow.service.v1.Environments/CheckUpgrade", + request_serializer=environments.CheckUpgradeRequest.serialize, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["check_upgrade"] + @property def create_user_workloads_secret( self, @@ -1021,6 +1052,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.check_upgrade: gapic_v1.method_async.wrap_method( + self.check_upgrade, + default_timeout=None, + client_info=client_info, + ), self.create_user_workloads_secret: gapic_v1.method_async.wrap_method( self.create_user_workloads_secret, default_timeout=None, diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/rest.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/rest.py index a735deff7f54..6a8be243f2a9 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/rest.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/services/environments/transports/rest.py @@ -72,6 +72,14 @@ class EnvironmentsRestInterceptor: .. code-block:: python class MyCustomEnvironmentsInterceptor(EnvironmentsRestInterceptor): + def pre_check_upgrade(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_check_upgrade(self, response): + logging.log(f"Received response: {response}") + return response + def pre_create_environment(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -254,6 +262,29 @@ def post_update_user_workloads_secret(self, response): """ + def pre_check_upgrade( + self, + request: environments.CheckUpgradeRequest, + metadata: Sequence[Tuple[str, str]], + ) -> Tuple[environments.CheckUpgradeRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for check_upgrade + + Override in a subclass to manipulate the request or metadata + before they are sent to the Environments server. + """ + return request, metadata + + def post_check_upgrade( + self, response: operations_pb2.Operation + ) -> operations_pb2.Operation: + """Post-rpc interceptor for check_upgrade + + Override in a subclass to manipulate the response + after it is returned by the Environments server but before + it is returned to user code. + """ + return response + def pre_create_environment( self, request: environments.CreateEnvironmentRequest, @@ -988,6 +1019,100 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: # Return the client from cache. return self._operations_client + class _CheckUpgrade(EnvironmentsRestStub): + def __hash__(self): + return hash("CheckUpgrade") + + __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 + } + + def __call__( + self, + request: environments.CheckUpgradeRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.Operation: + r"""Call the check upgrade method over HTTP. + + Args: + request (~.environments.CheckUpgradeRequest): + The request object. Request to check whether image + upgrade will succeed. + 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, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + ~.operations_pb2.Operation: + This resource represents a + long-running operation that is the + result of a network API call. + + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "post", + "uri": "/v1/{environment=projects/*/locations/*/environments/*}:checkUpgrade", + "body": "*", + }, + ] + request, metadata = self._interceptor.pre_check_upgrade(request, metadata) + pb_request = environments.CheckUpgradeRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + + # Jsonify the request body + + body = json_format.MessageToJson( + transcoded_request["body"], use_integers_for_enums=True + ) + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update(self._get_unset_required_fields(query_params)) + + query_params["$alt"] = "json;enum-encoding=int" + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=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_check_upgrade(resp) + return resp + class _CreateEnvironment(EnvironmentsRestStub): def __hash__(self): return hash("CreateEnvironment") @@ -2958,6 +3083,14 @@ def __call__( resp = self._interceptor.post_update_user_workloads_secret(resp) return resp + @property + def check_upgrade( + self, + ) -> Callable[[environments.CheckUpgradeRequest], 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._CheckUpgrade(self._session, self._host, self._interceptor) # type: ignore + @property def create_environment( self, diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/__init__.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/__init__.py index 873b24e7b536..05d6386b5e43 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/__init__.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/__init__.py @@ -14,6 +14,8 @@ # limitations under the License. # from .environments import ( + AirflowMetadataRetentionPolicyConfig, + CheckUpgradeRequest, CheckUpgradeResponse, CloudDataLineageIntegration, CreateEnvironmentRequest, @@ -81,6 +83,8 @@ from .operations import OperationMetadata __all__ = ( + "AirflowMetadataRetentionPolicyConfig", + "CheckUpgradeRequest", "CheckUpgradeResponse", "CloudDataLineageIntegration", "CreateEnvironmentRequest", diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/environments.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/environments.py index 5414a93ee0fe..05f5049857da 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/environments.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1/types/environments.py @@ -79,9 +79,11 @@ "MasterAuthorizedNetworksConfig", "CloudDataLineageIntegration", "Environment", + "CheckUpgradeRequest", "CheckUpgradeResponse", "DataRetentionConfig", "TaskLogsRetentionConfig", + "AirflowMetadataRetentionPolicyConfig", }, ) @@ -1392,8 +1394,10 @@ class EnvironmentConfig(proto.Message): hours. If this value is omitted, the default value for - maintenance window will be applied. The default - value is Saturday and Sunday 00-06 GMT. + maintenance window is applied. By default, + maintenance windows are from 00:00:00 to + 04:00:00 (GMT) on Friday, Saturday, and Sunday + every week. workloads_config (google.cloud.orchestration.airflow.service_v1.types.WorkloadsConfig): Optional. The workloads configuration settings for the GKE cluster associated with the Cloud Composer environment. The @@ -2559,6 +2563,9 @@ class TriggererResource(proto.Message): class DagProcessorResource(proto.Message): r"""Configuration for resources used by Airflow DAG processors. + This field is supported for Cloud Composer environments in versions + composer-3.\ *.*-airflow-*.*.\* and newer. + Attributes: cpu (float): Optional. CPU request and limit for a single @@ -2778,6 +2785,8 @@ class Environment(proto.Message): <= 128 bytes in size. satisfies_pzs (bool): Output only. Reserved for future use. + satisfies_pzi (bool): + Output only. Reserved for future use. storage_config (google.cloud.orchestration.airflow.service_v1.types.StorageConfig): Optional. Storage configuration for this environment. @@ -2850,6 +2859,10 @@ class State(proto.Enum): proto.BOOL, number=8, ) + satisfies_pzi: bool = proto.Field( + proto.BOOL, + number=10, + ) storage_config: "StorageConfig" = proto.Field( proto.MESSAGE, number=9, @@ -2857,6 +2870,58 @@ class State(proto.Enum): ) +class CheckUpgradeRequest(proto.Message): + r"""Request to check whether image upgrade will succeed. + + Attributes: + environment (str): + Required. The resource name of the + environment to check upgrade for, in the form: + + "projects/{projectId}/locations/{locationId}/environments/{environmentId}". + image_version (str): + Optional. The version of the software running in the + environment. This encapsulates both the version of Cloud + Composer functionality and the version of Apache Airflow. It + must match the regular expression + ``composer-([0-9]+(\.[0-9]+\.[0-9]+(-preview\.[0-9]+)?)?|latest)-airflow-([0-9]+(\.[0-9]+(\.[0-9]+)?)?)``. + When used as input, the server also checks if the provided + version is supported and denies the request for an + unsupported version. + + The Cloud Composer portion of the image version is a full + `semantic version `__, or an alias in + the form of major version number or ``latest``. When an + alias is provided, the server replaces it with the current + Cloud Composer version that satisfies the alias. + + The Apache Airflow portion of the image version is a full + semantic version that points to one of the supported Apache + Airflow versions, or an alias in the form of only major or + major.minor versions specified. When an alias is provided, + the server replaces it with the latest Apache Airflow + version that satisfies the alias and is supported in the + given Cloud Composer version. + + In all cases, the resolved image version is stored in the + same field. + + See also `version + list `__ + and `versioning + overview `__. + """ + + environment: str = proto.Field( + proto.STRING, + number=1, + ) + image_version: str = proto.Field( + proto.STRING, + number=2, + ) + + class CheckUpgradeResponse(proto.Message): r"""Message containing information about the result of an upgrade check operation. @@ -2927,11 +2992,21 @@ class DataRetentionConfig(proto.Message): mechanism. Attributes: + airflow_metadata_retention_config (google.cloud.orchestration.airflow.service_v1.types.AirflowMetadataRetentionPolicyConfig): + Optional. The retention policy for airflow + metadata database. task_logs_retention_config (google.cloud.orchestration.airflow.service_v1.types.TaskLogsRetentionConfig): Optional. The configuration settings for task logs retention """ + airflow_metadata_retention_config: "AirflowMetadataRetentionPolicyConfig" = ( + proto.Field( + proto.MESSAGE, + number=1, + message="AirflowMetadataRetentionPolicyConfig", + ) + ) task_logs_retention_config: "TaskLogsRetentionConfig" = proto.Field( proto.MESSAGE, number=2, @@ -2945,8 +3020,7 @@ class TaskLogsRetentionConfig(proto.Message): Attributes: storage_mode (google.cloud.orchestration.airflow.service_v1.types.TaskLogsRetentionConfig.TaskLogsStorageMode): Optional. The mode of storage for Airflow - workers task logs. For details, see - go/composer-store-task-logs-in-cloud-logging-only-design-doc + workers task logs. """ class TaskLogsStorageMode(proto.Enum): @@ -2973,4 +3047,43 @@ class TaskLogsStorageMode(proto.Enum): ) +class AirflowMetadataRetentionPolicyConfig(proto.Message): + r"""The policy for airflow metadata database retention. + + Attributes: + retention_mode (google.cloud.orchestration.airflow.service_v1.types.AirflowMetadataRetentionPolicyConfig.RetentionMode): + Optional. Retention can be either enabled or + disabled. + retention_days (int): + Optional. How many days data should be + retained for. + """ + + class RetentionMode(proto.Enum): + r"""Describes retention policy. + + Values: + RETENTION_MODE_UNSPECIFIED (0): + Default mode doesn't change environment + parameters. + RETENTION_MODE_ENABLED (1): + Retention policy is enabled. + RETENTION_MODE_DISABLED (2): + Retention policy is disabled. + """ + RETENTION_MODE_UNSPECIFIED = 0 + RETENTION_MODE_ENABLED = 1 + RETENTION_MODE_DISABLED = 2 + + retention_mode: RetentionMode = proto.Field( + proto.ENUM, + number=1, + enum=RetentionMode, + ) + retention_days: int = proto.Field( + proto.INT32, + number=2, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1beta1/types/environments.py b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1beta1/types/environments.py index b83e81eb209b..a7d029a3a945 100644 --- a/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1beta1/types/environments.py +++ b/packages/google-cloud-orchestration-airflow/google/cloud/orchestration/airflow/service_v1beta1/types/environments.py @@ -2957,6 +2957,8 @@ class Environment(proto.Message): <= 128 bytes in size. satisfies_pzs (bool): Output only. Reserved for future use. + satisfies_pzi (bool): + Output only. Reserved for future use. storage_config (google.cloud.orchestration.airflow.service_v1beta1.types.StorageConfig): Optional. Storage configuration for this environment. @@ -3029,6 +3031,10 @@ class State(proto.Enum): proto.BOOL, number=8, ) + satisfies_pzi: bool = proto.Field( + proto.BOOL, + number=10, + ) storage_config: "StorageConfig" = proto.Field( proto.MESSAGE, number=9, diff --git a/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_async.py b/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_async.py new file mode 100644 index 000000000000..6b0d1d7e863b --- /dev/null +++ b/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_async.py @@ -0,0 +1,56 @@ +# -*- 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 CheckUpgrade +# 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-orchestration-airflow + + +# [START composer_v1_generated_Environments_CheckUpgrade_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.orchestration.airflow import service_v1 + + +async def sample_check_upgrade(): + # Create a client + client = service_v1.EnvironmentsAsyncClient() + + # Initialize request argument(s) + request = service_v1.CheckUpgradeRequest( + environment="environment_value", + ) + + # Make the request + operation = client.check_upgrade(request=request) + + print("Waiting for operation to complete...") + + response = (await operation).result() + + # Handle the response + print(response) + +# [END composer_v1_generated_Environments_CheckUpgrade_async] diff --git a/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_sync.py b/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_sync.py new file mode 100644 index 000000000000..f8e30156ed47 --- /dev/null +++ b/packages/google-cloud-orchestration-airflow/samples/generated_samples/composer_v1_generated_environments_check_upgrade_sync.py @@ -0,0 +1,56 @@ +# -*- 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 CheckUpgrade +# 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-orchestration-airflow + + +# [START composer_v1_generated_Environments_CheckUpgrade_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.orchestration.airflow import service_v1 + + +def sample_check_upgrade(): + # Create a client + client = service_v1.EnvironmentsClient() + + # Initialize request argument(s) + request = service_v1.CheckUpgradeRequest( + environment="environment_value", + ) + + # Make the request + operation = client.check_upgrade(request=request) + + print("Waiting for operation to complete...") + + response = operation.result() + + # Handle the response + print(response) + +# [END composer_v1_generated_Environments_CheckUpgrade_sync] diff --git a/packages/google-cloud-orchestration-airflow/samples/generated_samples/snippet_metadata_google.cloud.orchestration.airflow.service.v1.json b/packages/google-cloud-orchestration-airflow/samples/generated_samples/snippet_metadata_google.cloud.orchestration.airflow.service.v1.json index 3fffc7af33c9..389370672713 100644 --- a/packages/google-cloud-orchestration-airflow/samples/generated_samples/snippet_metadata_google.cloud.orchestration.airflow.service.v1.json +++ b/packages/google-cloud-orchestration-airflow/samples/generated_samples/snippet_metadata_google.cloud.orchestration.airflow.service.v1.json @@ -11,6 +11,159 @@ "version": "0.1.0" }, "snippets": [ + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.orchestration.airflow.service_v1.EnvironmentsAsyncClient", + "shortName": "EnvironmentsAsyncClient" + }, + "fullName": "google.cloud.orchestration.airflow.service_v1.EnvironmentsAsyncClient.check_upgrade", + "method": { + "fullName": "google.cloud.orchestration.airflow.service.v1.Environments.CheckUpgrade", + "service": { + "fullName": "google.cloud.orchestration.airflow.service.v1.Environments", + "shortName": "Environments" + }, + "shortName": "CheckUpgrade" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.orchestration.airflow.service_v1.types.CheckUpgradeRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.api_core.operation_async.AsyncOperation", + "shortName": "check_upgrade" + }, + "description": "Sample for CheckUpgrade", + "file": "composer_v1_generated_environments_check_upgrade_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "composer_v1_generated_Environments_CheckUpgrade_async", + "segments": [ + { + "end": 55, + "start": 27, + "type": "FULL" + }, + { + "end": 55, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 52, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 56, + "start": 53, + "type": "RESPONSE_HANDLING" + } + ], + "title": "composer_v1_generated_environments_check_upgrade_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.orchestration.airflow.service_v1.EnvironmentsClient", + "shortName": "EnvironmentsClient" + }, + "fullName": "google.cloud.orchestration.airflow.service_v1.EnvironmentsClient.check_upgrade", + "method": { + "fullName": "google.cloud.orchestration.airflow.service.v1.Environments.CheckUpgrade", + "service": { + "fullName": "google.cloud.orchestration.airflow.service.v1.Environments", + "shortName": "Environments" + }, + "shortName": "CheckUpgrade" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.orchestration.airflow.service_v1.types.CheckUpgradeRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.api_core.operation.Operation", + "shortName": "check_upgrade" + }, + "description": "Sample for CheckUpgrade", + "file": "composer_v1_generated_environments_check_upgrade_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "composer_v1_generated_Environments_CheckUpgrade_sync", + "segments": [ + { + "end": 55, + "start": 27, + "type": "FULL" + }, + { + "end": 55, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 52, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 56, + "start": 53, + "type": "RESPONSE_HANDLING" + } + ], + "title": "composer_v1_generated_environments_check_upgrade_sync.py" + }, { "canonical": true, "clientMethod": { diff --git a/packages/google-cloud-orchestration-airflow/scripts/fixup_service_v1_keywords.py b/packages/google-cloud-orchestration-airflow/scripts/fixup_service_v1_keywords.py index dcf2905a8c0c..1d453d96293f 100644 --- a/packages/google-cloud-orchestration-airflow/scripts/fixup_service_v1_keywords.py +++ b/packages/google-cloud-orchestration-airflow/scripts/fixup_service_v1_keywords.py @@ -39,6 +39,7 @@ def partition( class serviceCallTransformer(cst.CSTTransformer): CTRL_PARAMS: Tuple[str] = ('retry', 'timeout', 'metadata') METHOD_TO_PARAMS: Dict[str, Tuple[str]] = { + 'check_upgrade': ('environment', 'image_version', ), 'create_environment': ('parent', 'environment', ), 'create_user_workloads_config_map': ('parent', 'user_workloads_config_map', ), 'create_user_workloads_secret': ('parent', 'user_workloads_secret', ), diff --git a/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1/test_environments.py b/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1/test_environments.py index 3751f10f0bbf..ddd117361ace 100644 --- a/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1/test_environments.py +++ b/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1/test_environments.py @@ -1531,6 +1531,7 @@ def test_get_environment(request_type, transport: str = "grpc"): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) response = client.get_environment(request) @@ -1546,6 +1547,7 @@ def test_get_environment(request_type, transport: str = "grpc"): assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True def test_get_environment_empty_call(): @@ -1648,6 +1650,7 @@ async def test_get_environment_empty_call_async(): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) ) response = await client.get_environment() @@ -1720,6 +1723,7 @@ async def test_get_environment_async( uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) ) response = await client.get_environment(request) @@ -1736,6 +1740,7 @@ async def test_get_environment_async( assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True @pytest.mark.asyncio @@ -4783,6 +4788,293 @@ async def test_list_workloads_async_pages(): assert page_.raw_page.next_page_token == token +@pytest.mark.parametrize( + "request_type", + [ + environments.CheckUpgradeRequest, + dict, + ], +) +def test_check_upgrade(request_type, transport: str = "grpc"): + client = EnvironmentsClient( + 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.check_upgrade), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation(name="operations/spam") + response = client.check_upgrade(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = environments.CheckUpgradeRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, future.Future) + + +def test_check_upgrade_empty_call(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = EnvironmentsClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.check_upgrade), "__call__") as call: + call.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client.check_upgrade() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == environments.CheckUpgradeRequest() + + +def test_check_upgrade_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 = EnvironmentsClient( + 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 = environments.CheckUpgradeRequest( + environment="environment_value", + image_version="image_version_value", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.check_upgrade), "__call__") as call: + call.return_value.name = ( + "foo" # operation_request.operation in compute client(s) expect a string. + ) + client.check_upgrade(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == environments.CheckUpgradeRequest( + environment="environment_value", + image_version="image_version_value", + ) + + +def test_check_upgrade_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 = EnvironmentsClient( + 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.check_upgrade 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.check_upgrade] = mock_rpc + request = {} + client.check_upgrade(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.check_upgrade(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_check_upgrade_empty_call_async(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = EnvironmentsAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.check_upgrade), "__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.check_upgrade() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == environments.CheckUpgradeRequest() + + +@pytest.mark.asyncio +async def test_check_upgrade_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 = EnvironmentsAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + 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.check_upgrade + 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.check_upgrade + ] = mock_rpc + + request = {} + await client.check_upgrade(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.check_upgrade(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_check_upgrade_async( + transport: str = "grpc_asyncio", request_type=environments.CheckUpgradeRequest +): + client = EnvironmentsAsyncClient( + 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.check_upgrade), "__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.check_upgrade(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = environments.CheckUpgradeRequest() + 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_check_upgrade_async_from_dict(): + await test_check_upgrade_async(request_type=dict) + + +def test_check_upgrade_field_headers(): + client = EnvironmentsClient( + 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 = environments.CheckUpgradeRequest() + + request.environment = "environment_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.check_upgrade), "__call__") as call: + call.return_value = operations_pb2.Operation(name="operations/op") + client.check_upgrade(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", + "environment=environment_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_check_upgrade_field_headers_async(): + client = EnvironmentsAsyncClient( + 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 = environments.CheckUpgradeRequest() + + request.environment = "environment_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.check_upgrade), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation(name="operations/op") + ) + await client.check_upgrade(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", + "environment=environment_value", + ) in kw["metadata"] + + @pytest.mark.parametrize( "request_type", [ @@ -10400,7 +10692,11 @@ def test_create_environment_rest(request_type): }, "resilience_mode": 1, "data_retention_config": { - "task_logs_retention_config": {"storage_mode": 1} + "airflow_metadata_retention_config": { + "retention_mode": 1, + "retention_days": 1512, + }, + "task_logs_retention_config": {"storage_mode": 1}, }, }, "uuid": "uuid_value", @@ -10409,6 +10705,7 @@ def test_create_environment_rest(request_type): "update_time": {}, "labels": {}, "satisfies_pzs": True, + "satisfies_pzi": True, "storage_config": {"bucket": "bucket_value"}, } # The version of a generated dependency at test runtime may differ from the version used during generation. @@ -10713,6 +11010,7 @@ def test_get_environment_rest(request_type): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) # Wrap the value into a proper Response obj @@ -10732,6 +11030,7 @@ def test_get_environment_rest(request_type): assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True def test_get_environment_rest_use_cached_wrapped_rpc(): @@ -11324,7 +11623,11 @@ def test_update_environment_rest(request_type): }, "resilience_mode": 1, "data_retention_config": { - "task_logs_retention_config": {"storage_mode": 1} + "airflow_metadata_retention_config": { + "retention_mode": 1, + "retention_days": 1512, + }, + "task_logs_retention_config": {"storage_mode": 1}, }, }, "uuid": "uuid_value", @@ -11333,6 +11636,7 @@ def test_update_environment_rest(request_type): "update_time": {}, "labels": {}, "satisfies_pzs": True, + "satisfies_pzi": True, "storage_config": {"bucket": "bucket_value"}, } # The version of a generated dependency at test runtime may differ from the version used during generation. @@ -12747,6 +13051,255 @@ def test_list_workloads_rest_pager(transport: str = "rest"): assert page_.raw_page.next_page_token == token +@pytest.mark.parametrize( + "request_type", + [ + environments.CheckUpgradeRequest, + dict, + ], +) +def test_check_upgrade_rest(request_type): + client = EnvironmentsClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # send a request that will satisfy transcoding + request_init = { + "environment": "projects/sample1/locations/sample2/environments/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 = operations_pb2.Operation(name="operations/spam") + + # 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 + response = client.check_upgrade(request) + + # Establish that the response is the type that we expect. + assert response.operation.name == "operations/spam" + + +def test_check_upgrade_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 = EnvironmentsClient( + 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.check_upgrade 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.check_upgrade] = mock_rpc + + request = {} + client.check_upgrade(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.check_upgrade(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_check_upgrade_rest_required_fields( + request_type=environments.CheckUpgradeRequest, +): + transport_class = transports.EnvironmentsRestTransport + + request_init = {} + request_init["environment"] = "" + 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() + ).check_upgrade._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["environment"] = "environment_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).check_upgrade._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "environment" in jsonified_request + assert jsonified_request["environment"] == "environment_value" + + client = EnvironmentsClient( + 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 + + response = client.check_upgrade(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_check_upgrade_rest_unset_required_fields(): + transport = transports.EnvironmentsRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.check_upgrade._get_unset_required_fields({}) + assert set(unset_fields) == (set(()) & set(("environment",))) + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_check_upgrade_rest_interceptors(null_interceptor): + transport = transports.EnvironmentsRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.EnvironmentsRestInterceptor(), + ) + client = EnvironmentsClient(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.EnvironmentsRestInterceptor, "post_check_upgrade" + ) as post, mock.patch.object( + transports.EnvironmentsRestInterceptor, "pre_check_upgrade" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + pb_message = environments.CheckUpgradeRequest.pb( + environments.CheckUpgradeRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = Response() + req.return_value.status_code = 200 + req.return_value.request = PreparedRequest() + req.return_value._content = json_format.MessageToJson( + operations_pb2.Operation() + ) + + request = environments.CheckUpgradeRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = operations_pb2.Operation() + + client.check_upgrade( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + + +def test_check_upgrade_rest_bad_request( + transport: str = "rest", request_type=environments.CheckUpgradeRequest +): + client = EnvironmentsClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # send a request that will satisfy transcoding + request_init = { + "environment": "projects/sample1/locations/sample2/environments/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 = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.check_upgrade(request) + + +def test_check_upgrade_rest_error(): + client = EnvironmentsClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + @pytest.mark.parametrize( "request_type", [ @@ -17095,6 +17648,7 @@ def test_environments_base_transport(): "stop_airflow_command", "poll_airflow_command", "list_workloads", + "check_upgrade", "create_user_workloads_secret", "get_user_workloads_secret", "list_user_workloads_secrets", @@ -17416,6 +17970,9 @@ def test_environments_client_transport_session_collision(transport_name): session1 = client1.transport.list_workloads._session session2 = client2.transport.list_workloads._session assert session1 != session2 + session1 = client1.transport.check_upgrade._session + session2 = client2.transport.check_upgrade._session + assert session1 != session2 session1 = client1.transport.create_user_workloads_secret._session session2 = client2.transport.create_user_workloads_secret._session assert session1 != session2 diff --git a/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1beta1/test_environments.py b/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1beta1/test_environments.py index 5cef23532261..053e5f28d823 100644 --- a/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1beta1/test_environments.py +++ b/packages/google-cloud-orchestration-airflow/tests/unit/gapic/service_v1beta1/test_environments.py @@ -1534,6 +1534,7 @@ def test_get_environment(request_type, transport: str = "grpc"): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) response = client.get_environment(request) @@ -1549,6 +1550,7 @@ def test_get_environment(request_type, transport: str = "grpc"): assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True def test_get_environment_empty_call(): @@ -1651,6 +1653,7 @@ async def test_get_environment_empty_call_async(): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) ) response = await client.get_environment() @@ -1723,6 +1726,7 @@ async def test_get_environment_async( uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) ) response = await client.get_environment(request) @@ -1739,6 +1743,7 @@ async def test_get_environment_async( assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True @pytest.mark.asyncio @@ -11008,6 +11013,7 @@ def test_create_environment_rest(request_type): "update_time": {}, "labels": {}, "satisfies_pzs": True, + "satisfies_pzi": True, "storage_config": {"bucket": "bucket_value"}, } # The version of a generated dependency at test runtime may differ from the version used during generation. @@ -11312,6 +11318,7 @@ def test_get_environment_rest(request_type): uuid="uuid_value", state=environments.Environment.State.CREATING, satisfies_pzs=True, + satisfies_pzi=True, ) # Wrap the value into a proper Response obj @@ -11331,6 +11338,7 @@ def test_get_environment_rest(request_type): assert response.uuid == "uuid_value" assert response.state == environments.Environment.State.CREATING assert response.satisfies_pzs is True + assert response.satisfies_pzi is True def test_get_environment_rest_use_cached_wrapped_rpc(): @@ -11938,6 +11946,7 @@ def test_update_environment_rest(request_type): "update_time": {}, "labels": {}, "satisfies_pzs": True, + "satisfies_pzi": True, "storage_config": {"bucket": "bucket_value"}, } # The version of a generated dependency at test runtime may differ from the version used during generation.