Source code for airflow.providers.google.cloud.hooks.cloud_composer
## Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you 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__importannotationsimportasyncioimporttimefromcollections.abcimportMutableSequence,SequencefromtypingimportTYPE_CHECKINGfromgoogle.api_core.client_optionsimportClientOptionsfromgoogle.api_core.gapic_v1.methodimportDEFAULT,_MethodDefaultfromgoogle.cloud.orchestration.airflow.service_v1import(EnvironmentsAsyncClient,EnvironmentsClient,ImageVersionsClient,PollAirflowCommandResponse,)fromairflow.exceptionsimportAirflowExceptionfromairflow.providers.google.common.constsimportCLIENT_INFOfromairflow.providers.google.common.hooks.base_googleimportGoogleBaseHookifTYPE_CHECKING:fromgoogle.api_core.operationimportOperationfromgoogle.api_core.operation_asyncimportAsyncOperationfromgoogle.api_core.retryimportRetryfromgoogle.api_core.retry_asyncimportAsyncRetryfromgoogle.cloud.orchestration.airflow.service_v1.services.environments.pagersimport(ListEnvironmentsPager,)fromgoogle.cloud.orchestration.airflow.service_v1.services.image_versions.pagersimport(ListImageVersionsPager,)fromgoogle.cloud.orchestration.airflow.service_v1.typesimport(Environment,ExecuteAirflowCommandResponse,)fromgoogle.protobuf.field_mask_pb2importFieldMask
[docs]classCloudComposerHook(GoogleBaseHook):"""Hook for Google Cloud Composer APIs."""
[docs]defget_environment_client(self)->EnvironmentsClient:"""Retrieve client library object that allow access Environments service."""returnEnvironmentsClient(credentials=self.get_credentials(),client_info=CLIENT_INFO,client_options=self.client_options,)
[docs]defget_image_versions_client(self)->ImageVersionsClient:"""Retrieve client library object that allow access Image Versions service."""returnImageVersionsClient(credentials=self.get_credentials(),client_info=CLIENT_INFO,client_options=self.client_options,)
[docs]defwait_for_operation(self,operation:Operation,timeout:float|None=None):"""Wait for long-lasting operation to complete."""try:returnoperation.result(timeout=timeout)exceptException:error=operation.exception(timeout=timeout)raiseAirflowException(error)
[docs]defcreate_environment(self,project_id:str,region:str,environment:Environment|dict,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->Operation:""" Create a new environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment: The environment to create. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()result=client.create_environment(request={"parent":self.get_parent(project_id,region),"environment":environment},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defdelete_environment(self,project_id:str,region:str,environment_id:str,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->Operation:""" Delete an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()name=self.get_environment_name(project_id,region,environment_id)result=client.delete_environment(request={"name":name},retry=retry,timeout=timeout,metadata=metadata)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_environment(self,project_id:str,region:str,environment_id:str,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->Environment:""" Get an existing environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()result=client.get_environment(request={"name":self.get_environment_name(project_id,region,environment_id)},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]deflist_environments(self,project_id:str,region:str,page_size:int|None=None,page_token:str|None=None,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->ListEnvironmentsPager:""" List environments. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param page_size: The maximum number of environments to return. :param page_token: The next_page_token value returned from a previous List request, if any. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()result=client.list_environments(request={"parent":self.get_parent(project_id,region),"page_size":page_size,"page_token":page_token,},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defupdate_environment(self,project_id:str,region:str,environment_id:str,environment:Environment|dict,update_mask:dict|FieldMask,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->Operation:r""" Update an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param environment: A patch environment. Fields specified by the ``updateMask`` will be copied from the patch environment into the environment under update. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param update_mask: Required. A comma-separated list of paths, relative to ``Environment``, of fields to update. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.protobuf.field_mask_pb2.FieldMask` :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()name=self.get_environment_name(project_id,region,environment_id)result=client.update_environment(request={"name":name,"environment":environment,"update_mask":update_mask},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]deflist_image_versions(self,project_id:str,region:str,page_size:int|None=None,page_token:str|None=None,include_past_releases:bool=False,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->ListImageVersionsPager:""" List ImageVersions for provided location. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param page_size: The maximum number of environments to return. :param page_token: The next_page_token value returned from a previous List request, if any. :param include_past_releases: Flag to include past releases :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_image_versions_client()result=client.list_image_versions(request={"parent":self.get_parent(project_id,region),"page_size":page_size,"page_token":page_token,"include_past_releases":include_past_releases,},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defexecute_airflow_command(self,project_id:str,region:str,environment_id:str,command:str,subcommand:str,parameters:MutableSequence[str],retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->ExecuteAirflowCommandResponse:""" Execute Airflow command for provided Composer environment. :param project_id: The ID of the Google Cloud project that the service belongs to. :param region: The ID of the Google Cloud region that the service belongs to. :param environment_id: The ID of the Google Cloud environment that the service belongs to. :param command: Airflow command. :param subcommand: Airflow subcommand. :param parameters: Parameters for the Airflow command/subcommand as an array of arguments. It may contain positional arguments like ``["my-dag-id"]``, key-value parameters like ``["--foo=bar"]`` or ``["--foo","bar"]``, or other flags like ``["-f"]``. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()result=client.execute_airflow_command(request={"environment":self.get_environment_name(project_id,region,environment_id),"command":command,"subcommand":subcommand,"parameters":parameters,},retry=retry,timeout=timeout,metadata=metadata,)returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defpoll_airflow_command(self,project_id:str,region:str,environment_id:str,execution_id:str,pod:str,pod_namespace:str,next_line_number:int,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->PollAirflowCommandResponse:""" Poll Airflow command execution result for provided Composer environment. :param project_id: The ID of the Google Cloud project that the service belongs to. :param region: The ID of the Google Cloud region that the service belongs to. :param environment_id: The ID of the Google Cloud environment that the service belongs to. :param execution_id: The unique ID of the command execution. :param pod: The name of the pod where the command is executed. :param pod_namespace: The namespace of the pod where the command is executed. :param next_line_number: Line number from which new logs should be fetched. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()result=client.poll_airflow_command(request={"environment":self.get_environment_name(project_id,region,environment_id),"execution_id":execution_id,"pod":pod,"pod_namespace":pod_namespace,"next_line_number":next_line_number,},retry=retry,timeout=timeout,metadata=metadata,)returnresult
[docs]defwait_command_execution_result(self,project_id:str,region:str,environment_id:str,execution_cmd_info:dict,retry:Retry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),poll_interval:int=10,)->dict:whileTrue:try:result=self.poll_airflow_command(project_id=project_id,region=region,environment_id=environment_id,execution_id=execution_cmd_info["execution_id"],pod=execution_cmd_info["pod"],pod_namespace=execution_cmd_info["pod_namespace"],next_line_number=1,retry=retry,timeout=timeout,metadata=metadata,)exceptExceptionasex:self.log.exception("Exception occurred while polling CMD result")raiseAirflowException(ex)result_dict=PollAirflowCommandResponse.to_dict(result)ifresult_dict["output_end"]:returnresult_dictself.log.info("Waiting for result...")time.sleep(poll_interval)
[docs]classCloudComposerAsyncHook(GoogleBaseHook):"""Hook for Google Cloud Composer async APIs."""
[docs]asyncdefcreate_environment(self,project_id:str,region:str,environment:Environment|dict,retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->AsyncOperation:""" Create a new environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment: The environment to create. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()returnawaitclient.create_environment(request={"parent":self.get_parent(project_id,region),"environment":environment},retry=retry,timeout=timeout,metadata=metadata,)
@GoogleBaseHook.fallback_to_default_project_id
[docs]asyncdefdelete_environment(self,project_id:str,region:str,environment_id:str,retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->AsyncOperation:""" Delete an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()name=self.get_environment_name(project_id,region,environment_id)returnawaitclient.delete_environment(request={"name":name},retry=retry,timeout=timeout,metadata=metadata)
@GoogleBaseHook.fallback_to_default_project_id
[docs]asyncdefupdate_environment(self,project_id:str,region:str,environment_id:str,environment:Environment|dict,update_mask:dict|FieldMask,retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->AsyncOperation:r""" Update an environment. :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to. :param environment: A patch environment. Fields specified by the ``updateMask`` will be copied from the patch environment into the environment under update. This corresponds to the ``environment`` field on the ``request`` instance; if ``request`` is provided, this should not be set. :param update_mask: Required. A comma-separated list of paths, relative to ``Environment``, of fields to update. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.protobuf.field_mask_pb2.FieldMask` :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()name=self.get_environment_name(project_id,region,environment_id)returnawaitclient.update_environment(request={"name":name,"environment":environment,"update_mask":update_mask},retry=retry,timeout=timeout,metadata=metadata,)
@GoogleBaseHook.fallback_to_default_project_id
[docs]asyncdefexecute_airflow_command(self,project_id:str,region:str,environment_id:str,command:str,subcommand:str,parameters:MutableSequence[str],retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->AsyncOperation:""" Execute Airflow command for provided Composer environment. :param project_id: The ID of the Google Cloud project that the service belongs to. :param region: The ID of the Google Cloud region that the service belongs to. :param environment_id: The ID of the Google Cloud environment that the service belongs to. :param command: Airflow command. :param subcommand: Airflow subcommand. :param parameters: Parameters for the Airflow command/subcommand as an array of arguments. It may contain positional arguments like ``["my-dag-id"]``, key-value parameters like ``["--foo=bar"]`` or ``["--foo","bar"]``, or other flags like ``["-f"]``. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()returnawaitclient.execute_airflow_command(request={"environment":self.get_environment_name(project_id,region,environment_id),"command":command,"subcommand":subcommand,"parameters":parameters,},retry=retry,timeout=timeout,metadata=metadata,)
@GoogleBaseHook.fallback_to_default_project_id
[docs]asyncdefpoll_airflow_command(self,project_id:str,region:str,environment_id:str,execution_id:str,pod:str,pod_namespace:str,next_line_number:int,retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),)->AsyncOperation:""" Poll Airflow command execution result for provided Composer environment. :param project_id: The ID of the Google Cloud project that the service belongs to. :param region: The ID of the Google Cloud region that the service belongs to. :param environment_id: The ID of the Google Cloud environment that the service belongs to. :param execution_id: The unique ID of the command execution. :param pod: The name of the pod where the command is executed. :param pod_namespace: The namespace of the pod where the command is executed. :param next_line_number: Line number from which new logs should be fetched. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. """client=self.get_environment_client()returnawaitclient.poll_airflow_command(request={"environment":self.get_environment_name(project_id,region,environment_id),"execution_id":execution_id,"pod":pod,"pod_namespace":pod_namespace,"next_line_number":next_line_number,},retry=retry,timeout=timeout,metadata=metadata,)
[docs]asyncdefwait_command_execution_result(self,project_id:str,region:str,environment_id:str,execution_cmd_info:dict,retry:AsyncRetry|_MethodDefault=DEFAULT,timeout:float|None=None,metadata:Sequence[tuple[str,str]]=(),poll_interval:int=10,)->dict:whileTrue:try:result=awaitself.poll_airflow_command(project_id=project_id,region=region,environment_id=environment_id,execution_id=execution_cmd_info["execution_id"],pod=execution_cmd_info["pod"],pod_namespace=execution_cmd_info["pod_namespace"],next_line_number=1,retry=retry,timeout=timeout,metadata=metadata,)exceptExceptionasex:self.log.exception("Exception occurred while polling CMD result")raiseAirflowException(ex)result_dict=PollAirflowCommandResponse.to_dict(result)ifresult_dict["output_end"]:returnresult_dictself.log.info("Sleeping for %s seconds.",poll_interval)awaitasyncio.sleep(poll_interval)