apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/transfers/postgres_to_gcs.py

@@ -31,7 +31,7 @@ import pendulum
 from slugify import slugify
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
-from airflow.providers.postgres.hooks.postgres import PostgresHook
+from airflow.providers.postgres.hooks.postgres import USE_PSYCOPG3, PostgresHook
 
 if TYPE_CHECKING:
     from airflow.providers.openlineage.extractors import OperatorLineage
@@ -52,14 +52,24 @@ class _PostgresServerSideCursorDecorator:
         self.initialized = False
 
     def __iter__(self):
+        """Make the cursor iterable."""
         return self
 
     def __next__(self):
+        """Fetch next row from the cursor."""
+        if USE_PSYCOPG3:
+            if self.rows:
+                return self.rows.pop()
+            self.initialized = True
+            row = self.cursor.fetchone()
+            if row is None:
+                raise StopIteration
+            return row
+        # psycopg2
         if self.rows:
             return self.rows.pop()
-        else:
-            self.initialized = True
-            return next(self.cursor)
+        self.initialized = True
+        return next(self.cursor)
 
     @property
     def description(self):
@@ -142,13 +152,29 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
         return PostgresHook(postgres_conn_id=self.postgres_conn_id)
 
     def query(self):
-        """Query Postgres and returns a cursor to the results."""
+        """Execute the query and return a cursor."""
         conn = self.db_hook.get_conn()
-        cursor = conn.cursor(name=self._unique_name())
-        cursor.execute(self.sql, self.parameters)
-        if self.use_server_side_cursor:
-            cursor.itersize = self.cursor_itersize
-            return _PostgresServerSideCursorDecorator(cursor)
+
+        if USE_PSYCOPG3:
+            from psycopg.types.json import register_default_adapters
+
+            # Register JSON handlers for this connection if not already done
+            register_default_adapters(conn)
+
+            if self.use_server_side_cursor:
+                cursor_name = f"airflow_{self.task_id.replace('-', '_')}_{uuid.uuid4().hex}"[:63]
+                cursor = conn.cursor(name=cursor_name)
+                cursor.itersize = self.cursor_itersize
+                cursor.execute(self.sql, self.parameters)
+                return _PostgresServerSideCursorDecorator(cursor)
+            cursor = conn.cursor()
+            cursor.execute(self.sql, self.parameters)
+        else:
+            cursor = conn.cursor(name=self._unique_name())
+            cursor.execute(self.sql, self.parameters)
+            if self.use_server_side_cursor:
+                cursor.itersize = self.cursor_itersize
+                return _PostgresServerSideCursorDecorator(cursor)
         return cursor
 
     def field_to_bigquery(self, field) -> dict[str, str]:
@@ -183,8 +209,14 @@ class PostgresToGCSOperator(BaseSQLToGCSOperator):
                 hours=formatted_time.tm_hour, minutes=formatted_time.tm_min, seconds=formatted_time.tm_sec
             )
             return str(time_delta)
-        if stringify_dict and isinstance(value, dict):
-            return json.dumps(value)
+        if stringify_dict:
+            if USE_PSYCOPG3:
+                from psycopg.types.json import Json
+
+                if isinstance(value, (dict, Json)):
+                    return json.dumps(value)
+            elif isinstance(value, dict):
+                return json.dumps(value)
         if isinstance(value, Decimal):
             return float(value)
         return value

airflow/providers/google/cloud/transfers/s3_to_gcs.py

@@ -57,7 +57,7 @@ except ImportError:
     from airflow.providers.amazon.aws.operators.s3_list import S3ListOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class S3ToGCSOperator(S3ListOperator):
@@ -181,21 +181,27 @@ class S3ToGCSOperator(S3ListOperator):
                 'The destination Google Cloud Storage path must end with a slash "/" or be empty.'
             )
 
-    def execute(self, context: Context):
-        self._check_inputs()
+    def _get_files(self, context: Context, gcs_hook: GCSHook) -> list[str]:
         # use the super method to list all the files in an S3 bucket/key
         s3_objects = super().execute(context)
 
+        if not self.replace:
+            s3_objects = self.exclude_existing_objects(s3_objects=s3_objects, gcs_hook=gcs_hook)
+
+        return s3_objects
+
+    def execute(self, context: Context):
+        self._check_inputs()
         gcs_hook = GCSHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.google_impersonation_chain,
         )
-        if not self.replace:
-            s3_objects = self.exclude_existing_objects(s3_objects=s3_objects, gcs_hook=gcs_hook)
-
         s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
+
+        s3_objects = self._get_files(context, gcs_hook)
         if not s3_objects:
             self.log.info("In sync, no files needed to be uploaded to Google Cloud Storage")
+
         elif self.deferrable:
             self.transfer_files_async(s3_objects, gcs_hook, s3_hook)
         else:

airflow/providers/google/cloud/transfers/salesforce_to_gcs.py

@@ -21,12 +21,12 @@ import tempfile
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 from airflow.providers.salesforce.hooks.salesforce import SalesforceHook
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class SalesforceToGcsOperator(BaseOperator):

airflow/providers/google/cloud/transfers/sftp_to_gcs.py

@@ -26,12 +26,12 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 from airflow.providers.sftp.hooks.sftp import SFTPHook
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 WILDCARD = "*"
@@ -73,6 +73,13 @@ class SFTPToGCSOperator(BaseOperator):
         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).
     :param sftp_prefetch: Whether to enable SFTP prefetch, the default is True.
+    :param use_stream: Determines the transfer method from SFTP to GCS.
+        When ``False`` (default), the file downloads locally
+        then uploads (may require significant disk space).
+        When ``True``, the file streams directly without using local disk.
+        Defaults to ``False``.
+    :param fail_on_file_not_exist: If True, operator fails when file does not exist,
+        if False, operator will not fail and skips transfer. Default is True.
     """
 
     template_fields: Sequence[str] = (
@@ -95,6 +102,8 @@ class SFTPToGCSOperator(BaseOperator):
         move_object: bool = False,
         impersonation_chain: str | Sequence[str] | None = None,
         sftp_prefetch: bool = True,
+        use_stream: bool = False,
+        fail_on_file_not_exist: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -109,6 +118,8 @@ class SFTPToGCSOperator(BaseOperator):
         self.move_object = move_object
         self.impersonation_chain = impersonation_chain
         self.sftp_prefetch = sftp_prefetch
+        self.use_stream = use_stream
+        self.fail_on_file_not_exist = fail_on_file_not_exist
 
     @cached_property
     def sftp_hook(self):
@@ -149,7 +160,13 @@ class SFTPToGCSOperator(BaseOperator):
             destination_object = (
                 self.destination_path if self.destination_path else self.source_path.rsplit("/", 1)[1]
             )
-            self._copy_single_object(gcs_hook, self.sftp_hook, self.source_path, destination_object)
+            try:
+                self._copy_single_object(gcs_hook, self.sftp_hook, self.source_path, destination_object)
+            except FileNotFoundError as e:
+                if self.fail_on_file_not_exist:
+                    raise e
+                self.log.info("File %s not found on SFTP server. Skipping transfer.", self.source_path)
+                return
 
     def _copy_single_object(
         self,
@@ -165,17 +182,22 @@ class SFTPToGCSOperator(BaseOperator):
             self.destination_bucket,
             destination_object,
         )
-
-        with NamedTemporaryFile("w") as tmp:
-            sftp_hook.retrieve_file(source_path, tmp.name, prefetch=self.sftp_prefetch)
-
-            gcs_hook.upload(
-                bucket_name=self.destination_bucket,
-                object_name=destination_object,
-                filename=tmp.name,
-                mime_type=self.mime_type,
-                gzip=self.gzip,
-            )
+        if self.use_stream:
+            dest_bucket = gcs_hook.get_bucket(self.destination_bucket)
+            dest_blob = dest_bucket.blob(destination_object)
+            with dest_blob.open("wb") as write_stream:
+                sftp_hook.retrieve_file(source_path, write_stream, prefetch=self.sftp_prefetch)
+        else:
+            with NamedTemporaryFile("w") as tmp:
+                sftp_hook.retrieve_file(source_path, tmp.name, prefetch=self.sftp_prefetch)
+
+                gcs_hook.upload(
+                    bucket_name=self.destination_bucket,
+                    object_name=destination_object,
+                    filename=tmp.name,
+                    mime_type=self.mime_type,
+                    gzip=self.gzip,
+                )
 
         if self.move_object:
             self.log.info("Executing delete of %s", source_path)

airflow/providers/google/cloud/transfers/sheets_to_gcs.py

@@ -21,12 +21,12 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING, Any
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.suite.hooks.sheets import GSheetsHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class GoogleSheetsToGCSOperator(BaseOperator):
@@ -130,5 +130,5 @@ class GoogleSheetsToGCSOperator(BaseOperator):
             gcs_path_to_file = self._upload_data(gcs_hook, sheet_hook, sheet_range, data)
             destination_array.append(gcs_path_to_file)
 
-        self.xcom_push(context, "destination_objects", destination_array)
+        context["ti"].xcom_push(key="destination_objects", value=destination_array)
         return destination_array

airflow/providers/google/cloud/transfers/sql_to_gcs.py

@@ -30,12 +30,12 @@ from typing import TYPE_CHECKING, Any
 import pyarrow as pa
 import pyarrow.parquet as pq
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.providers.common.compat.openlineage.facet import OutputDataset
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BaseSQLToGCSOperator(BaseOperator):
@@ -295,21 +295,21 @@ class BaseSQLToGCSOperator(BaseOperator):
 
             # Proceed to write the row to the localfile
             if self.export_format == "csv":
-                row = self.convert_types(schema, col_type_dict, row)
+                row2 = self.convert_types(schema, col_type_dict, row)
                 if self.null_marker is not None:
-                    row = [value or self.null_marker for value in row]
-                csv_writer.writerow(row)
+                    row2 = [value or self.null_marker for value in row2]
+                csv_writer.writerow(row2)
             elif self.export_format == "parquet":
-                row = self.convert_types(schema, col_type_dict, row)
+                row2 = self.convert_types(schema, col_type_dict, row)
                 if self.null_marker is not None:
-                    row = [value or self.null_marker for value in row]
-                rows_buffer.append(row)
+                    row2 = [value or self.null_marker for value in row2]
+                rows_buffer.append(row2)
                 if len(rows_buffer) >= self.parquet_row_group_size:
                     self._write_rows_to_parquet(parquet_writer, rows_buffer)
                     rows_buffer = []
             else:
-                row = self.convert_types(schema, col_type_dict, row)
-                row_dict = dict(zip(schema, row))
+                row2 = self.convert_types(schema, col_type_dict, row)
+                row_dict = dict(zip(schema, row2))
 
                 json.dump(row_dict, tmp_file_handle, sort_keys=True, ensure_ascii=False)
 

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

@@ -22,17 +22,21 @@ from typing import TYPE_CHECKING, Any, SupportsAbs
 
 from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientResponseError
+from asgiref.sync import sync_to_async
 
 from airflow.exceptions import AirflowException
-from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryAsyncHook, BigQueryTableAsyncHook
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils.session import provide_session
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
 
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.models.taskinstance import TaskInstance
+    from airflow.utils.session import provide_session
+
 
 class BigQueryInsertJobTrigger(BaseTrigger):
     """
@@ -97,16 +101,39 @@ class BigQueryInsertJobTrigger(BaseTrigger):
             },
         )
 
-    @provide_session
-    def get_task_instance(self, session: Session) -> TaskInstance:
-        query = session.query(TaskInstance).filter(
-            TaskInstance.dag_id == self.task_instance.dag_id,
-            TaskInstance.task_id == self.task_instance.task_id,
-            TaskInstance.run_id == self.task_instance.run_id,
-            TaskInstance.map_index == self.task_instance.map_index,
+    if not AIRFLOW_V_3_0_PLUS:
+
+        @provide_session
+        def get_task_instance(self, session: Session) -> TaskInstance:
+            query = session.query(TaskInstance).filter(
+                TaskInstance.dag_id == self.task_instance.dag_id,
+                TaskInstance.task_id == self.task_instance.task_id,
+                TaskInstance.run_id == self.task_instance.run_id,
+                TaskInstance.map_index == self.task_instance.map_index,
+            )
+            task_instance = query.one_or_none()
+            if task_instance is None:
+                raise AirflowException(
+                    "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
+                    self.task_instance.dag_id,
+                    self.task_instance.task_id,
+                    self.task_instance.run_id,
+                    self.task_instance.map_index,
+                )
+            return task_instance
+
+    async def get_task_state(self):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+            dag_id=self.task_instance.dag_id,
+            task_ids=[self.task_instance.task_id],
+            run_ids=[self.task_instance.run_id],
+            map_index=self.task_instance.map_index,
         )
-        task_instance = query.one_or_none()
-        if task_instance is None:
+        try:
+            task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
+        except Exception:
             raise AirflowException(
                 "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
                 self.task_instance.dag_id,
@@ -114,20 +141,24 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                 self.task_instance.run_id,
                 self.task_instance.map_index,
             )
-        return task_instance
+        return task_state
 
-    def safe_to_cancel(self) -> bool:
+    async def safe_to_cancel(self) -> bool:
         """
         Whether it is safe to cancel the external job which is being executed by this trigger.
 
         This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
         Because in those cases, we should NOT cancel the external job.
         """
-        # Database query is needed to get the latest state of the task instance.
-        task_instance = self.get_task_instance()  # type: ignore[call-arg]
-        return task_instance.state != TaskInstanceState.DEFERRED
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+        if AIRFLOW_V_3_0_PLUS:
+            task_state = await self.get_task_state()
+        else:
+            # Database query is needed to get the latest state of the task instance.
+            task_instance = self.get_task_instance()  # type: ignore[call-arg]
+            task_state = task_instance.state
+        return task_state != TaskInstanceState.DEFERRED
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:
@@ -136,6 +167,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                     job_id=self.job_id, project_id=self.project_id, location=self.location
                 )
                 if job_status["status"] == "success":
+                    self.log.info("BigQuery Job succeeded")
                     yield TriggerEvent(
                         {
                             "job_id": self.job_id,
@@ -145,7 +177,13 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                     )
                     return
                 elif job_status["status"] == "error":
-                    yield TriggerEvent(job_status)
+                    self.log.info("BigQuery Job failed: %s", job_status)
+                    yield TriggerEvent(
+                        {
+                            "status": job_status["status"],
+                            "message": job_status["message"],
+                        }
+                    )
                     return
                 else:
                     self.log.info(
@@ -155,7 +193,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                     )
                     await asyncio.sleep(self.poll_interval)
         except asyncio.CancelledError:
-            if self.job_id and self.cancel_on_kill and self.safe_to_cancel():
+            if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
                 self.log.info(
                     "The job is safe to cancel the as airflow TaskInstance is not in deferred state."
                 )
@@ -165,9 +203,7 @@ class BigQueryInsertJobTrigger(BaseTrigger):
                     self.location,
                     self.job_id,
                 )
-                await hook.cancel_job(  # type: ignore[union-attr]
-                    job_id=self.job_id, project_id=self.project_id, location=self.location
-                )
+                await hook.cancel_job(job_id=self.job_id, project_id=self.project_id, location=self.location)
             else:
                 self.log.info(
                     "Trigger may have shutdown. Skipping to cancel job because the airflow "
@@ -204,7 +240,7 @@ class BigQueryCheckTrigger(BigQueryInsertJobTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:
@@ -281,7 +317,7 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent with response data."""
         hook = self._get_async_hook()
         try:
@@ -305,7 +341,12 @@ class BigQueryGetDataTrigger(BigQueryInsertJobTrigger):
                     )
                     return
                 elif job_status["status"] == "error":
-                    yield TriggerEvent(job_status)
+                    yield TriggerEvent(
+                        {
+                            "status": job_status["status"],
+                            "message": job_status["message"],
+                        }
+                    )
                     return
                 else:
                     self.log.info(
@@ -406,7 +447,7 @@ class BigQueryIntervalCheckTrigger(BigQueryInsertJobTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:
@@ -554,7 +595,7 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:
@@ -564,9 +605,9 @@ class BigQueryValueCheckTrigger(BigQueryInsertJobTrigger):
                 if response_from_hook["status"] == "success":
                     query_results = await hook.get_job_output(job_id=self.job_id, project_id=self.project_id)
                     records = hook.get_records(query_results)
-                    records = records.pop(0) if records else None
-                    hook.value_check(self.sql, self.pass_value, records, self.tolerance)
-                    yield TriggerEvent({"status": "success", "message": "Job completed", "records": records})
+                    _records = records.pop(0) if records else None
+                    hook.value_check(self.sql, self.pass_value, _records, self.tolerance)
+                    yield TriggerEvent({"status": "success", "message": "Job completed", "records": _records})
                     return
                 elif response_from_hook["status"] == "pending":
                     self.log.info("Query is still running...")
@@ -640,7 +681,7 @@ class BigQueryTableExistenceTrigger(BaseTrigger):
             gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Will run until the table exists in the Google Big Query."""
         try:
             while True:
@@ -723,7 +764,7 @@ class BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Will run until the table exists in the Google Big Query."""
         hook = BigQueryAsyncHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         job_id = None
@@ -744,7 +785,7 @@ class BigQueryTablePartitionExistenceTrigger(BigQueryTableExistenceTrigger):
                         return
                     job_id = None
                 elif job_status["status"] == "error":
-                    yield TriggerEvent(job_status)
+                    yield TriggerEvent({"status": job_status["status"]})
                     return
                 self.log.info("Sleeping for %s seconds.", self.poll_interval)
                 await asyncio.sleep(self.poll_interval)

airflow/providers/google/cloud/triggers/bigquery_dts.py

@@ -21,9 +21,10 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
+from google.cloud.bigquery_datatransfer_v1 import TransferRun, TransferState
+
 from airflow.providers.google.cloud.hooks.bigquery_dts import AsyncBiqQueryDataTransferServiceHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from google.cloud.bigquery_datatransfer_v1 import TransferRun, TransferState
 
 
 class BigQueryDataTransferRunTrigger(BaseTrigger):

airflow/providers/google/cloud/triggers/cloud_batch.py

@@ -20,9 +20,10 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
+from google.cloud.batch_v1 import Job, JobStatus
+
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from google.cloud.batch_v1 import Job, JobStatus
 
 DEFAULT_BATCH_LOCATION = "us-central1"
 

airflow/providers/google/cloud/triggers/cloud_build.py

@@ -20,9 +20,10 @@ import asyncio
 from collections.abc import AsyncIterator, Sequence
 from typing import Any
 
+from google.cloud.devtools.cloudbuild_v1.types import Build
+
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildAsyncHook
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from google.cloud.devtools.cloudbuild_v1.types import Build
 
 
 class CloudBuildCreateBuildTrigger(BaseTrigger):
@@ -75,7 +76,7 @@ class CloudBuildCreateBuildTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current build execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try: