apache-airflow-providers-google 16.0.0a1__py3-none-any.whl → 16.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/workflows.py

@@ -147,7 +147,6 @@ class WorkflowsCreateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -235,7 +234,6 @@ class WorkflowsUpdateWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -368,7 +366,6 @@ class WorkflowsListWorkflowsOperator(GoogleCloudBaseOperator):
 
         WorkflowsListOfWorkflowsLink.persist(
             context=context,
-            task_instance=self,
             project_id=self.project_id or hook.project_id,
         )
 
@@ -434,7 +431,6 @@ class WorkflowsGetWorkflowOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -505,11 +501,10 @@ class WorkflowsCreateExecutionOperator(GoogleCloudBaseOperator):
             metadata=self.metadata,
         )
         execution_id = execution.name.split("/")[-1]
-        self.xcom_push(context, key="execution_id", value=execution_id)
+        context["task_instance"].xcom_push(key="execution_id", value=execution_id)
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=execution_id,
@@ -582,7 +577,6 @@ class WorkflowsCancelExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,
@@ -661,7 +655,6 @@ class WorkflowsListExecutionsOperator(GoogleCloudBaseOperator):
 
         WorkflowsWorkflowDetailsLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             project_id=self.project_id or hook.project_id,
@@ -737,7 +730,6 @@ class WorkflowsGetExecutionOperator(GoogleCloudBaseOperator):
 
         WorkflowsExecutionLink.persist(
             context=context,
-            task_instance=self,
             location_id=self.location,
             workflow_id=self.workflow_id,
             execution_id=self.execution_id,

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

@@ -31,7 +31,7 @@ from airflow.providers.google.cloud.triggers.bigquery import (
     BigQueryTableExistenceTrigger,
     BigQueryTablePartitionExistenceTrigger,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -28,7 +28,12 @@ from google.cloud.bigquery_datatransfer_v1 import TransferState
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

airflow/providers/google/cloud/sensors/bigtable.py

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import google.api_core.exceptions
 from google.cloud.bigtable import enums
@@ -30,7 +30,12 @@ from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import BigtableTablesLink
 from airflow.providers.google.cloud.operators.bigtable import BigtableValidationMixin
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -89,6 +94,13 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
         self.impersonation_chain = impersonation_chain
         super().__init__(**kwargs)
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "instance_id": self.instance_id,
+            "project_id": self.project_id,
+        }
+
     def poke(self, context: Context) -> bool:
         hook = BigtableHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -119,5 +131,5 @@ class BigtableTableReplicationCompletedSensor(BaseSensorOperator, BigtableValida
             return False
 
         self.log.info("Table '%s' is replicated.", self.table_id)
-        BigtableTablesLink.persist(context=context, task_instance=self)
+        BigtableTablesLink.persist(context=context)
         return True

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

@@ -33,7 +33,12 @@ from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
 from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerDAGRunTrigger
 from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:

airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py

@@ -35,7 +35,12 @@ from airflow.providers.google.cloud.triggers.cloud_storage_transfer_service impo
     CloudStorageTransferServiceCheckJobStatusTrigger,
 )
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/dataflow.py

@@ -19,9 +19,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -37,7 +37,7 @@ from airflow.providers.google.cloud.triggers.dataflow import (
     DataflowJobStatusTrigger,
 )
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/dataform.py

@@ -24,7 +24,12 @@ from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataform import DataformHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/datafusion.py

@@ -25,7 +25,12 @@ from typing import TYPE_CHECKING
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/dataplex.py

@@ -37,7 +37,12 @@ from airflow.providers.google.cloud.hooks.dataplex import (
     AirflowDataQualityScanResultTimeoutException,
     DataplexHook,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 
 class TaskState:

airflow/providers/google/cloud/sensors/dataprep.py

@@ -23,7 +23,12 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/dataproc.py

@@ -29,7 +29,12 @@ from google.cloud.dataproc_v1.types import Batch, JobStatus
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc import DataprocHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/dataproc_metastore.py

@@ -23,7 +23,12 @@ from typing import TYPE_CHECKING
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
 from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from google.api_core.operation import Operation

airflow/providers/google/cloud/sensors/gcs.py

@@ -21,9 +21,9 @@ from __future__ import annotations
 
 import os
 import textwrap
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from google.cloud.storage.retry import DEFAULT_RETRY
 
@@ -36,7 +36,13 @@ from airflow.providers.google.cloud.triggers.gcs import (
     GCSPrefixBlobTrigger,
     GCSUploadSessionTrigger,
 )
-from airflow.sensors.base import BaseSensorOperator, poke_mode_only
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+    from airflow.sdk.bases.sensor import poke_mode_only
+else:
+    from airflow.sensors.base import BaseSensorOperator, poke_mode_only  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

airflow/providers/google/cloud/sensors/looker.py

@@ -23,7 +23,12 @@ from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/pubsub.py

@@ -19,9 +19,9 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1.types import ReceivedMessage
@@ -30,7 +30,12 @@ from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.pubsub import PubSubHook
 from airflow.providers.google.cloud.triggers.pubsub import PubsubPullTrigger
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/tasks.py

@@ -24,7 +24,12 @@ from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/vertex_ai/feature_store.py

@@ -25,7 +25,12 @@ from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/cloud/sensors/workflows.py

@@ -25,7 +25,12 @@ from google.cloud.workflows.executions_v1beta import Execution
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -21,8 +21,8 @@ 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
 
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook

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

@@ -23,8 +23,8 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url, gcs_object_is_directory
+from airflow.providers.google.version_compat import BaseOperator
 
 try:
     from airflow.providers.microsoft.azure.hooks.fileshare import AzureFileShareHook

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

@@ -22,9 +22,10 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -73,6 +74,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
         If set as a sequence, the identities from the list must grant
         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 project_id: Google Cloud Project where the job is running
     """
 
     template_fields: Sequence[str] = (
@@ -93,6 +95,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
         write_disposition: str = "WRITE_EMPTY",
         create_disposition: str = "CREATE_IF_NEEDED",
         gcp_conn_id: str = "google_cloud_default",
+        project_id: str = PROVIDE_PROJECT_ID,
         labels: dict | None = None,
         encryption_configuration: dict | None = None,
         location: str | None = None,
@@ -112,6 +115,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
         self.impersonation_chain = impersonation_chain
         self.hook: BigQueryHook | None = None
         self._job_conf: dict = {}
+        self.project_id = project_id
 
     def _prepare_job_configuration(self):
         self.source_project_dataset_tables = (
@@ -124,7 +128,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
         for source_project_dataset_table in self.source_project_dataset_tables:
             source_project, source_dataset, source_table = self.hook.split_tablename(
                 table_input=source_project_dataset_table,
-                default_project_id=self.hook.project_id,
+                default_project_id=self.project_id,
                 var_name="source_project_dataset_table",
             )
             source_project_dataset_tables_fixup.append(
@@ -133,7 +137,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
 
         destination_project, destination_dataset, destination_table = self.hook.split_tablename(
             table_input=self.destination_project_dataset_table,
-            default_project_id=self.hook.project_id,
+            default_project_id=self.project_id,
         )
         configuration = {
             "copy": {
@@ -168,18 +172,17 @@ class BigQueryToBigQueryOperator(BaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
 
-        if not self.hook.project_id:
-            raise ValueError("The project_id should be set")
+        if not self.project_id:
+            self.project_id = self.hook.project_id
 
         configuration = self._prepare_job_configuration()
         self._job_conf = self.hook.insert_job(
-            configuration=configuration, project_id=self.hook.project_id
+            configuration=configuration, project_id=self.project_id
         ).to_api_repr()
 
         dest_table_info = self._job_conf["configuration"]["copy"]["destinationTable"]
         BigQueryTableLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dest_table_info["datasetId"],
             project_id=dest_table_info["projectId"],
             table_id=dest_table_info["tableId"],

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

@@ -27,11 +27,11 @@ from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import BaseOperator
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
@@ -254,7 +254,6 @@ class BigQueryToGCSOperator(BaseOperator):
         dataset_id, project_id, table_id = conf["datasetId"], conf["projectId"], conf["tableId"]
         BigQueryTableLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
             table_id=table_id,

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

@@ -101,7 +101,6 @@ class BigQueryToMsSqlOperator(BigQueryToSqlBaseOperator):
         project_id, dataset_id, table_id = self.source_project_dataset_table.split(".")
         BigQueryTableLink.persist(
             context=context,
-            task_instance=self,
             dataset_id=dataset_id,
             project_id=project_id,
             table_id=table_id,

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

@@ -23,9 +23,9 @@ import abc
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
 from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.providers.common.sql.hooks.sql import DbApiHook

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

@@ -21,9 +21,9 @@ 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.calendar import GoogleCalendarHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from datetime import datetime

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

@@ -31,9 +31,9 @@ from uuid import UUID
 from cassandra.util import Date, OrderedMapSerializedKey, SortedSet, Time
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -26,9 +26,9 @@ from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.facebook.ads.hooks.ads import FacebookAdsReportingHook
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from facebook_business.adobjects.adsinsights import AdsInsights
@@ -208,7 +208,7 @@ class FacebookAdsReportToGcsOperator(BaseOperator):
 
     def _flush_rows(self, converted_rows: list[Any] | None, object_name: str):
         if converted_rows:
-            headers = converted_rows[0].keys()
+            headers = self.fields
             with tempfile.NamedTemporaryFile("w", suffix=".csv") as csvfile:
                 writer = csv.DictWriter(csvfile, fieldnames=headers)
                 writer.writeheader()

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

@@ -38,12 +38,12 @@ from google.cloud.bigquery.table import EncryptionConfiguration, Table, TableRef
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook, BigQueryJob
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink
 from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import BaseOperator
 from airflow.utils.helpers import merge_dicts
 
 if TYPE_CHECKING:
@@ -373,7 +373,6 @@ class GCSToBigQueryOperator(BaseOperator):
 
             BigQueryTableLink.persist(
                 context=context,
-                task_instance=self,
                 dataset_id=table_obj_api_repr["tableReference"]["datasetId"],
                 project_id=table_obj_api_repr["tableReference"]["projectId"],
                 table_id=table_obj_api_repr["tableReference"]["tableId"],

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

@@ -24,8 +24,8 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 WILDCARD = "*"
 

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

@@ -20,9 +20,9 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.models.xcom import MAX_XCOM_SIZE
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -26,8 +26,8 @@ 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
 
 WILDCARD = "*"

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

@@ -19,9 +19,9 @@ from __future__ import annotations
 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.suite.hooks.drive import GoogleDriveHook
+from airflow.providers.google.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -99,3 +99,7 @@ class GoogleDriveToGCSOperator(BaseOperator):
             bucket_name=self.bucket_name, object_name=self.object_name
         ) as file:
             gdrive_hook.download_file(file_id=file_metadata["id"], file_handle=file)
+
+    def dry_run(self):
+        """Perform a dry run of the operator."""
+        return None