diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 97375f5011fbe..0f70ff62fe756 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -117,7 +117,8 @@ def get_operation(self, operation_name: str, project_id: Optional[str] = None) - :return: The new, updated operation from Google Cloud """ return self.get_conn().get_operation( - project_id=project_id or self.project_id, zone=self.location, operation_id=operation_name + name=f'projects/{project_id or self.project_id}' + + f'/locations/{self.location}/operations/{operation_name}' ) @staticmethod @@ -170,11 +171,13 @@ def delete_cluster( :type timeout: float :return: The full url to the delete operation if successful, else None """ - self.log.info("Deleting (project_id=%s, zone=%s, cluster_id=%s)", project_id, self.location, name) + self.log.info("Deleting (project_id=%s, location=%s, cluster_id=%s)", project_id, self.location, name) try: resource = self.get_conn().delete_cluster( - project_id=project_id, zone=self.location, cluster_id=name, retry=retry, timeout=timeout + name=f'projects/{project_id}/locations/{self.location}/clusters/{name}', + retry=retry, + timeout=timeout, ) resource = self.wait_for_operation(resource) # Returns server-defined url for the resource @@ -223,11 +226,14 @@ def create_cluster( self._append_label(cluster, 'airflow-version', 'v' + version.version) self.log.info( - "Creating (project_id=%s, zone=%s, cluster_name=%s)", project_id, self.location, cluster.name + "Creating (project_id=%s, location=%s, cluster_name=%s)", project_id, self.location, cluster.name ) try: resource = self.get_conn().create_cluster( - project_id=project_id, zone=self.location, cluster=cluster, retry=retry, timeout=timeout + parent=f'projects/{project_id}/locations/{self.location}', + cluster=cluster, + retry=retry, + timeout=timeout, ) resource = self.wait_for_operation(resource) @@ -261,7 +267,7 @@ def get_cluster( :return: google.cloud.container_v1.types.Cluster """ self.log.info( - "Fetching cluster (project_id=%s, zone=%s, cluster_name=%s)", + "Fetching cluster (project_id=%s, location=%s, cluster_name=%s)", project_id or self.project_id, self.location, name, @@ -270,7 +276,9 @@ def get_cluster( return ( self.get_conn() .get_cluster( - project_id=project_id, zone=self.location, cluster_id=name, retry=retry, timeout=timeout + name=f'projects/{project_id}/locations/{self.location}/clusters/{name}', + retry=retry, + timeout=timeout, ) .self_link ) diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index 50ba3f9af747a..1b507fd9e6353 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -59,7 +59,7 @@ class GKEDeleteClusterOperator(BaseOperator): :type project_id: str :param name: The name of the resource to delete, in this case cluster name :type name: str - :param location: The name of the Google Compute Engine zone in which the cluster + :param location: The name of the Google Compute Engine zone or region in which the cluster resides. :type location: str :param gcp_conn_id: The connection ID to use connecting to Google Cloud. @@ -158,7 +158,7 @@ class GKECreateClusterOperator(BaseOperator): :param project_id: The Google Developers Console [project ID or project number] :type project_id: str - :param location: The name of the Google Compute Engine zone in which the cluster + :param location: The name of the Google Compute Engine or region in which the cluster resides. :type location: str :param body: The Cluster definition to create, can be protobuf or python dict, if @@ -273,13 +273,14 @@ class GKEStartPodOperator(KubernetesPodOperator): For more information on how to use this operator, take a look at the guide: :ref:`howto/operator:GKEStartPodOperator` - :param location: The name of the Google Kubernetes Engine zone in which the + :param location: The name of the Google Kubernetes Engine zone or region in which the cluster resides, e.g. 'us-central1-a' :type location: str :param cluster_name: The name of the Google Kubernetes Engine cluster the pod should be spawned in :type cluster_name: str :param use_internal_ip: Use the internal IP address as the endpoint. + :type use_internal_ip: bool :param project_id: The Google Developers Console project id :type project_id: str :param gcp_conn_id: The google cloud connection id to use. This allows for @@ -294,6 +295,8 @@ class GKEStartPodOperator(KubernetesPodOperator): Service Account Token Creator IAM role to the directly preceding identity, with first account from the list granting this role to the originating account (templated). :type impersonation_chain: Union[str, Sequence[str]] + :param regional: The location param is region name. + :type regional: bool """ template_fields = {'project_id', 'location', 'cluster_name'} | set(KubernetesPodOperator.template_fields) @@ -307,6 +310,7 @@ def __init__( project_id: Optional[str] = None, gcp_conn_id: str = 'google_cloud_default', impersonation_chain: Optional[Union[str, Sequence[str]]] = None, + regional: bool = False, **kwargs, ) -> None: super().__init__(**kwargs) @@ -316,6 +320,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id self.use_internal_ip = use_internal_ip self.impersonation_chain = impersonation_chain + self.regional = regional if self.gcp_conn_id is None: raise AirflowException( @@ -356,8 +361,6 @@ def execute(self, context) -> Optional[str]: "clusters", "get-credentials", self.cluster_name, - "--zone", - self.location, "--project", self.project_id, ] @@ -377,6 +380,11 @@ def execute(self, context) -> Optional[str]: 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) diff --git a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py index f42d7932fed68..641fcab0c74b8 100644 --- a/tests/providers/google/cloud/hooks/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/hooks/test_kubernetes_engine.py @@ -74,9 +74,7 @@ def test_delete_cluster(self, wait_mock, convert_mock, mock_project_id): ) client_delete.assert_called_once_with( - project_id=TEST_GCP_PROJECT_ID, - zone=GKE_ZONE, - cluster_id=CLUSTER_NAME, + name=f'projects/{TEST_GCP_PROJECT_ID}/locations/{GKE_ZONE}/clusters/{CLUSTER_NAME}', retry=retry_mock, timeout=timeout_mock, ) @@ -145,8 +143,7 @@ def test_create_cluster_proto(self, wait_mock, convert_mock, mock_project_id): ) client_create.assert_called_once_with( - project_id=TEST_GCP_PROJECT_ID, - zone=GKE_ZONE, + parent=f'projects/{TEST_GCP_PROJECT_ID}/locations/{GKE_ZONE}', cluster=mock_cluster_proto, retry=retry_mock, timeout=timeout_mock, @@ -173,8 +170,7 @@ def test_create_cluster_dict(self, wait_mock, convert_mock, mock_project_id): ) client_create.assert_called_once_with( - project_id=TEST_GCP_PROJECT_ID, - zone=GKE_ZONE, + parent=f'projects/{TEST_GCP_PROJECT_ID}/locations/{GKE_ZONE}', cluster=proto_mock, retry=retry_mock, timeout=timeout_mock, @@ -228,9 +224,7 @@ def test_get_cluster(self): ) client_get.assert_called_once_with( - project_id=TEST_GCP_PROJECT_ID, - zone=GKE_ZONE, - cluster_id=CLUSTER_NAME, + name=f'projects/{TEST_GCP_PROJECT_ID}/locations/{GKE_ZONE}/clusters/{CLUSTER_NAME}', retry=retry_mock, timeout=timeout_mock, ) @@ -256,7 +250,7 @@ def test_get_operation(self): self.gke_hook._client.get_operation = mock.Mock() self.gke_hook.get_operation('TEST_OP', project_id=TEST_GCP_PROJECT_ID) self.gke_hook._client.get_operation.assert_called_once_with( - project_id=TEST_GCP_PROJECT_ID, zone=GKE_ZONE, operation_id='TEST_OP' + name=f'projects/{TEST_GCP_PROJECT_ID}/locations/{GKE_ZONE}/operations/TEST_OP' ) def test_append_label(self): diff --git a/tests/providers/google/cloud/operators/test_kubernetes_engine.py b/tests/providers/google/cloud/operators/test_kubernetes_engine.py index 04df7b36682f2..58bd380ca176c 100644 --- a/tests/providers/google/cloud/operators/test_kubernetes_engine.py +++ b/tests/providers/google/cloud/operators/test_kubernetes_engine.py @@ -226,10 +226,53 @@ def test_execute(self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exe 'clusters', 'get-credentials', CLUSTER_NAME, + '--project', + TEST_GCP_PROJECT_ID, '--zone', PROJECT_LOCATION, + ] + ) + + assert self.gke_op.config_file == FILE_NAME + + @mock.patch.dict(os.environ, {}) + @mock.patch( + "airflow.hooks.base.BaseHook.get_connections", + return_value=[ + Connection( + extra=json.dumps( + {"extra__google_cloud_platform__keyfile_dict": '{"private_key": "r4nd0m_k3y"}'} + ) + ) + ], + ) + @mock.patch('airflow.providers.cncf.kubernetes.operators.kubernetes_pod.KubernetesPodOperator.execute') + @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.GoogleBaseHook') + @mock.patch('airflow.providers.google.cloud.operators.kubernetes_engine.execute_in_subprocess') + @mock.patch('tempfile.NamedTemporaryFile') + def test_execute_regional( + self, file_mock, mock_execute_in_subprocess, mock_gcp_hook, exec_mock, get_con_mock + ): + self.gke_op.regional = True + type(file_mock.return_value.__enter__.return_value).name = PropertyMock( + side_effect=[FILE_NAME, '/path/to/new-file'] + ) + + self.gke_op.execute(None) + + mock_gcp_hook.return_value.provide_authorized_gcloud.assert_called_once() + + mock_execute_in_subprocess.assert_called_once_with( + [ + 'gcloud', + 'container', + 'clusters', + 'get-credentials', + CLUSTER_NAME, '--project', TEST_GCP_PROJECT_ID, + '--region', + PROJECT_LOCATION, ] ) @@ -282,10 +325,10 @@ def test_execute_with_internal_ip( 'clusters', 'get-credentials', CLUSTER_NAME, - '--zone', - PROJECT_LOCATION, '--project', TEST_GCP_PROJECT_ID, + '--zone', + PROJECT_LOCATION, '--internal-ip', ] ) @@ -325,12 +368,12 @@ def test_execute_with_impersonation_service_account( 'clusters', 'get-credentials', CLUSTER_NAME, - '--zone', - PROJECT_LOCATION, '--project', TEST_GCP_PROJECT_ID, '--impersonate-service-account', 'test_account@example.com', + '--zone', + PROJECT_LOCATION, ] ) @@ -369,12 +412,12 @@ def test_execute_with_impersonation_service_chain_one_element( 'clusters', 'get-credentials', CLUSTER_NAME, - '--zone', - PROJECT_LOCATION, '--project', TEST_GCP_PROJECT_ID, '--impersonate-service-account', 'test_account@example.com', + '--zone', + PROJECT_LOCATION, ] )