def execute(self, context): hook = GoogleBaseHook(gcp_conn_id=self.gcp_conn_id) self.project_id = self.project_id or hook.project_id if not self.project_id: raise AirflowException( "The project id must be passed either as " "keyword project_id parameter or as project_id extra " "in GCP connection definition. Both are not set!") # Write config to a temp file and set the environment variable to point to it. # This is to avoid race conditions of reading/writing a single file with tempfile.NamedTemporaryFile() as conf_file,\ patch_environ({KUBE_CONFIG_ENV_VAR: conf_file.name}), \ hook.provide_authorized_gcloud(): # Attempt to get/update credentials # We call gcloud directly instead of using google-cloud-python api # because there is no way to write kubernetes config to a file, which is # required by KubernetesPodOperator. # The gcloud command looks at the env variable `KUBECONFIG` for where to save # the kubernetes config file. cmd = [ "gcloud", "container", "clusters", "get-credentials", self.cluster_name, "--zone", self.location, "--project", self.project_id ] if self.use_internal_ip: cmd.append('--internal-ip') execute_in_subprocess(cmd) # Tell `KubernetesPodOperator` where the config file is located self.config_file = os.environ[KUBE_CONFIG_ENV_VAR] return super().execute(context)
def execute(self, context: 'Context') -> Optional[str]: hook = GoogleBaseHook(gcp_conn_id=self.gcp_conn_id) self.project_id = self.project_id or hook.project_id if not self.project_id: raise AirflowException( "The project id must be passed either as " "keyword project_id parameter or as project_id extra " "in Google Cloud connection definition. Both are not set!") # Write config to a temp file and set the environment variable to point to it. # This is to avoid race conditions of reading/writing a single file with tempfile.NamedTemporaryFile() as conf_file, patch_environ( {KUBE_CONFIG_ENV_VAR: conf_file.name}), hook.provide_authorized_gcloud(): # Attempt to get/update credentials # We call gcloud directly instead of using google-cloud-python api # because there is no way to write kubernetes config to a file, which is # required by KubernetesPodOperator. # The gcloud command looks at the env variable `KUBECONFIG` for where to save # the kubernetes config file. cmd = [ "gcloud", "container", "clusters", "get-credentials", self.cluster_name, "--project", self.project_id, ] if self.impersonation_chain: if isinstance(self.impersonation_chain, str): impersonation_account = self.impersonation_chain elif len(self.impersonation_chain) == 1: impersonation_account = self.impersonation_chain[0] else: raise AirflowException( "Chained list of accounts is not supported, please specify only one service account" ) cmd.extend([ '--impersonate-service-account', impersonation_account, ]) if self.regional: cmd.append('--region') else: cmd.append('--zone') cmd.append(self.location) if self.use_internal_ip: cmd.append('--internal-ip') execute_in_subprocess(cmd) # Tell `KubernetesPodOperator` where the config file is located self.config_file = os.environ[KUBE_CONFIG_ENV_VAR] return super().execute(context)