apache-airflow-providers-google 10.17.0rc1__py3-none-any.whl → 10.18.0rc1__py3-none-any.whl

airflow/providers/google/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "10.17.0"
+__version__ = "10.18.0"
 
 try:
     from airflow import __version__ as airflow_version
@@ -35,8 +35,8 @@ except ImportError:
     from airflow.version import version as airflow_version
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.6.0"
+    "2.7.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.6.0+"
+        f"The package `apache-airflow-providers-google:{__version__}` needs Apache Airflow 2.7.0+"
     )

airflow/providers/google/cloud/hooks/automl.py

@@ -529,7 +529,7 @@ class CloudAutoMLHook(GoogleBaseHook):
         self,
         dataset_id: str,
         location: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         filter_: str | None = None,
         page_size: int | None = None,
         retry: Retry | _MethodDefault = DEFAULT,

airflow/providers/google/cloud/hooks/bigquery.py

@@ -59,7 +59,12 @@ from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.google.cloud.utils.bigquery import bq_cast
 from airflow.providers.google.cloud.utils.credentials_provider import _get_scopes
 from airflow.providers.google.common.consts import CLIENT_INFO
-from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook, get_field
+from airflow.providers.google.common.hooks.base_google import (
+    PROVIDE_PROJECT_ID,
+    GoogleBaseAsyncHook,
+    GoogleBaseHook,
+    get_field,
+)
 
 try:
     from airflow.utils.hashlib_wrapper import md5
@@ -198,7 +203,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         http_authorized = self._authorize()
         return build("bigquery", "v2", http=http_authorized, cache_discovery=False)
 
-    def get_client(self, project_id: str | None = None, location: str | None = None) -> Client:
+    def get_client(self, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None) -> Client:
         """Get an authenticated BigQuery Client.
 
         :param project_id: Project ID for the project which the client acts on behalf of.
@@ -250,7 +255,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     @staticmethod
     def _resolve_table_reference(
         table_resource: dict[str, Any],
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         dataset_id: str | None = None,
         table_id: str | None = None,
     ) -> dict[str, Any]:
@@ -360,7 +365,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     @GoogleBaseHook.fallback_to_default_project_id
     def create_empty_table(
         self,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         dataset_id: str | None = None,
         table_id: str | None = None,
         table_resource: dict[str, Any] | None = None,
@@ -474,7 +479,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def create_empty_dataset(
         self,
         dataset_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
         dataset_reference: dict[str, Any] | None = None,
         exists_ok: bool = True,
@@ -536,7 +541,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def get_dataset_tables(
         self,
         dataset_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         max_results: int | None = None,
         retry: Retry = DEFAULT_RETRY,
     ) -> list[dict[str, Any]]:
@@ -565,7 +570,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def delete_dataset(
         self,
         dataset_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         delete_contents: bool = False,
         retry: Retry = DEFAULT_RETRY,
     ) -> None:
@@ -614,7 +619,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         description: str | None = None,
         encryption_configuration: dict | None = None,
         location: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> Table:
         """Create an external table in the dataset with data from Google Cloud Storage.
 
@@ -750,7 +755,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         fields: list[str] | None = None,
         dataset_id: str | None = None,
         table_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> dict[str, Any]:
         """Change some fields of a table.
 
@@ -796,7 +801,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         self,
         dataset_id: str,
         table_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         description: str | None = None,
         expiration_time: int | None = None,
         external_data_configuration: dict | None = None,
@@ -953,7 +958,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         fields: Sequence[str],
         dataset_resource: dict[str, Any],
         dataset_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry = DEFAULT_RETRY,
     ) -> Dataset:
         """Change some fields of a dataset.
@@ -999,7 +1004,9 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         ),
         category=AirflowProviderDeprecationWarning,
     )
-    def patch_dataset(self, dataset_id: str, dataset_resource: dict, project_id: str | None = None) -> dict:
+    def patch_dataset(
+        self, dataset_id: str, dataset_resource: dict, project_id: str = PROVIDE_PROJECT_ID
+    ) -> dict:
         """Patches information in an existing dataset.
 
         It only replaces fields that are provided in the submitted dataset resource.
@@ -1047,7 +1054,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def get_dataset_tables_list(
         self,
         dataset_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         table_prefix: str | None = None,
         max_results: int | None = None,
     ) -> list[dict[str, Any]]:
@@ -1084,7 +1091,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     @GoogleBaseHook.fallback_to_default_project_id
     def get_datasets_list(
         self,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         include_all: bool = False,
         filter_: str | None = None,
         max_results: int | None = None,
@@ -1134,7 +1141,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         return datasets_list
 
     @GoogleBaseHook.fallback_to_default_project_id
-    def get_dataset(self, dataset_id: str, project_id: str | None = None) -> Dataset:
+    def get_dataset(self, dataset_id: str, project_id: str = PROVIDE_PROJECT_ID) -> Dataset:
         """Fetch the dataset referenced by *dataset_id*.
 
         :param dataset_id: The BigQuery Dataset ID
@@ -1158,7 +1165,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         view_dataset: str,
         view_table: str,
         view_project: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> dict[str, Any]:
         """Grant authorized view access of a dataset to a view table.
 
@@ -1210,7 +1217,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
 
     @GoogleBaseHook.fallback_to_default_project_id
     def run_table_upsert(
-        self, dataset_id: str, table_resource: dict[str, Any], project_id: str | None = None
+        self, dataset_id: str, table_resource: dict[str, Any], project_id: str = PROVIDE_PROJECT_ID
     ) -> dict[str, Any]:
         """Update a table if it exists, otherwise create a new one.
 
@@ -1267,7 +1274,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         self,
         table_id: str,
         not_found_ok: bool = True,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> None:
         """Delete an existing table from the dataset.
 
@@ -1334,7 +1341,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         selected_fields: list[str] | str | None = None,
         page_token: str | None = None,
         start_index: int | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
         retry: Retry = DEFAULT_RETRY,
         return_iterator: bool = False,
@@ -1387,7 +1394,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         return list(iterator)
 
     @GoogleBaseHook.fallback_to_default_project_id
-    def get_schema(self, dataset_id: str, table_id: str, project_id: str | None = None) -> dict:
+    def get_schema(self, dataset_id: str, table_id: str, project_id: str = PROVIDE_PROJECT_ID) -> dict:
         """Get the schema for a given dataset and table.
 
         .. seealso:: https://cloud.google.com/bigquery/docs/reference/v2/tables#resource
@@ -1409,7 +1416,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         include_policy_tags: bool,
         dataset_id: str,
         table_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> dict[str, Any]:
         """Update fields within a schema for a given dataset and table.
 
@@ -1502,7 +1509,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def poll_job_complete(
         self,
         job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
         retry: Retry = DEFAULT_RETRY,
     ) -> bool:
@@ -1532,7 +1539,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
     def cancel_job(
         self,
         job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
     ) -> None:
         """Cancel a job and wait for cancellation to complete.
@@ -1573,10 +1580,11 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
                 time.sleep(5)
 
     @GoogleBaseHook.fallback_to_default_project_id
+    @GoogleBaseHook.refresh_credentials_retry()
     def get_job(
         self,
         job_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
     ) -> CopyJob | QueryJob | LoadJob | ExtractJob | UnknownJob:
         """Retrieve a BigQuery job.
@@ -1607,7 +1615,7 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         self,
         configuration: dict,
         job_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         location: str | None = None,
         nowait: bool = False,
         retry: Retry = DEFAULT_RETRY,
@@ -2849,11 +2857,10 @@ class BigQueryCursor(BigQueryBaseCursor):
                 return None
 
             query_results = self._get_query_result()
-            if "rows" in query_results and query_results["rows"]:
+            if rows := query_results.get("rows"):
                 self.page_token = query_results.get("pageToken")
                 fields = query_results["schema"]["fields"]
                 col_types = [field["type"] for field in fields]
-                rows = query_results["rows"]
 
                 for dict_row in rows:
                     typed_row = [bq_cast(vs["v"], col_types[idx]) for idx, vs in enumerate(dict_row["f"])]
@@ -3305,7 +3312,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
         )
 
     async def _get_job(
-        self, job_id: str | None, project_id: str | None = None, location: str | None = None
+        self, job_id: str | None, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None
     ) -> CopyJob | QueryJob | LoadJob | ExtractJob | UnknownJob:
         """
         Get BigQuery job by its ID, project ID and location.
@@ -3348,7 +3355,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
         return hook.get_job(job_id=job_id, project_id=project_id, location=location)
 
     async def get_job_status(
-        self, job_id: str | None, project_id: str | None = None, location: str | None = None
+        self, job_id: str | None, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None
     ) -> dict[str, str]:
         job = await self._get_job(job_id=job_id, project_id=project_id, location=location)
         if job.state == "DONE":
@@ -3360,7 +3367,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
     async def get_job_output(
         self,
         job_id: str | None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ) -> dict[str, Any]:
         """Get the BigQuery job output for a given job ID asynchronously."""
         async with ClientSession() as session:
@@ -3373,7 +3380,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
         self,
         dataset_id: str | None,
         table_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
     ):
         """Create a new job and get the job_id using gcloud-aio."""
         async with ClientSession() as session:
@@ -3389,6 +3396,31 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
             job_query_resp = await job_client.query(query_request, cast(Session, session))
             return job_query_resp["jobReference"]["jobId"]
 
+    async def cancel_job(self, job_id: str, project_id: str | None, location: str | None) -> None:
+        """
+        Cancel a BigQuery job.
+
+        :param job_id: ID of the job to cancel.
+        :param project_id: Google Cloud Project where the job was running.
+        :param location: Location where the job was running.
+        """
+        async with ClientSession() as session:
+            token = await self.get_token(session=session)
+            job = Job(job_id=job_id, project=project_id, location=location, token=token, session=session)  # type: ignore[arg-type]
+
+            self.log.info(
+                "Attempting to cancel BigQuery job: %s in project: %s, location: %s",
+                job_id,
+                project_id,
+                location,
+            )
+            try:
+                await job.cancel()
+                self.log.info("Job %s cancellation requested.", job_id)
+            except Exception as e:
+                self.log.error("Failed to cancel BigQuery job %s: %s", job_id, str(e))
+                raise
+
     def get_records(self, query_results: dict[str, Any], as_dict: bool = False) -> list[Any]:
         """Convert a response from BigQuery to records.
 
@@ -3396,8 +3428,7 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
         :param as_dict: if True returns the result as a list of dictionaries, otherwise as list of lists.
         """
         buffer: list[Any] = []
-        if "rows" in query_results and query_results["rows"]:
-            rows = query_results["rows"]
+        if rows := query_results.get("rows"):
             fields = query_results["schema"]["fields"]
             fields_names = [field["name"] for field in fields]
             col_types = [field["type"] for field in fields]

airflow/providers/google/cloud/hooks/cloud_composer.py

@@ -17,7 +17,9 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+import asyncio
+import time
+from typing import TYPE_CHECKING, MutableSequence, Sequence
 
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
@@ -25,6 +27,7 @@ from google.cloud.orchestration.airflow.service_v1 import (
     EnvironmentsAsyncClient,
     EnvironmentsClient,
     ImageVersionsClient,
+    PollAirflowCommandResponse,
 )
 
 from airflow.exceptions import AirflowException
@@ -42,7 +45,10 @@ if TYPE_CHECKING:
     from google.cloud.orchestration.airflow.service_v1.services.image_versions.pagers import (
         ListImageVersionsPager,
     )
-    from google.cloud.orchestration.airflow.service_v1.types import Environment
+    from google.cloud.orchestration.airflow.service_v1.types import (
+        Environment,
+        ExecuteAirflowCommandResponse,
+    )
     from google.protobuf.field_mask_pb2 import FieldMask
 
 
@@ -294,6 +300,127 @@ class CloudComposerHook(GoogleBaseHook):
         )
         return result
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    def execute_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,
+        )
+        return result
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def poll_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,
+        )
+        return result
+
+    def wait_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:
+        while True:
+            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,
+                )
+            except Exception as ex:
+                self.log.exception("Exception occurred while polling CMD result")
+                raise AirflowException(ex)
+
+            result_dict = PollAirflowCommandResponse.to_dict(result)
+            if result_dict["output_end"]:
+                return result_dict
+
+            self.log.info("Waiting for result...")
+            time.sleep(poll_interval)
+
 
 class CloudComposerAsyncHook(GoogleBaseHook):
     """Hook for Google Cloud Composer async APIs."""
@@ -421,3 +548,124 @@ class CloudComposerAsyncHook(GoogleBaseHook):
             timeout=timeout,
             metadata=metadata,
         )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def execute_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()
+
+        return await 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,
+        )
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def poll_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()
+
+        return await 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,
+        )
+
+    async def wait_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:
+        while True:
+            try:
+                result = await 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,
+                )
+            except Exception as ex:
+                self.log.exception("Exception occurred while polling CMD result")
+                raise AirflowException(ex)
+
+            result_dict = PollAirflowCommandResponse.to_dict(result)
+            if result_dict["output_end"]:
+                return result_dict
+
+            self.log.info("Sleeping for %s seconds.", poll_interval)
+            await asyncio.sleep(poll_interval)