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

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

@@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     COUNTERS,
     METADATA,
@@ -35,10 +36,9 @@ 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
@@ -98,6 +98,7 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
         self.deferrable = deferrable
 
     def poke(self, context: Context) -> bool:
+        ti = context["ti"]
         hook = CloudDataTransferServiceHook(
             gcp_conn_id=self.gcp_cloud_conn_id,
             impersonation_chain=self.impersonation_chain,
@@ -113,13 +114,12 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
             operations=operations, expected_statuses=self.expected_statuses
         )
         if check:
-            self.xcom_push(key="sensed_operations", value=operations, context=context)
+            ti.xcom_push(key="sensed_operations", value=operations)
 
         project_id = self.project_id or hook.project_id
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=self.job_name,
             )
@@ -154,4 +154,5 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
         if event["status"] == "error":
             raise AirflowException(event["message"])
 
-        self.xcom_push(key="sensed_operations", value=event["operations"], context=context)
+        ti = context["ti"]
+        ti.xcom_push(key="sensed_operations", value=event["operations"])

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

@@ -19,12 +19,13 @@
 
 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
+from airflow.providers.common.compat.sdk import BaseSensorOperator, PokeReturnValue
 from airflow.providers.google.cloud.hooks.dataflow import (
     DEFAULT_DATAFLOW_LOCATION,
     DataflowHook,
@@ -37,10 +38,9 @@ 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class DataflowJobStatusSensor(BaseSensorOperator):
@@ -117,7 +117,7 @@ class DataflowJobStatusSensor(BaseSensorOperator):
 
         if job_status in self.expected_statuses:
             return True
-        elif job_status in DataflowJobStatus.TERMINAL_STATES:
+        if job_status in DataflowJobStatus.TERMINAL_STATES:
             message = f"Job with id '{self.job_id}' is already in terminal state: {job_status}"
             raise AirflowException(message)
 
@@ -342,7 +342,7 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
         self.deferrable = deferrable
         self.poll_interval = poll_interval
 
-    def poke(self, context: Context) -> bool:
+    def poke(self, context: Context) -> PokeReturnValue | bool:
         if self.fail_on_terminal_state:
             job = self.hook.get_job(
                 job_id=self.job_id,
@@ -359,8 +359,17 @@ class DataflowJobMessagesSensor(BaseSensorOperator):
             project_id=self.project_id,
             location=self.location,
         )
+        result = result if self.callback is None else self.callback(result)
+
+        if isinstance(result, PokeReturnValue):
+            return result
 
-        return result if self.callback is None else self.callback(result)
+        if bool(result):
+            return PokeReturnValue(
+                is_done=True,
+                xcom_value=result,
+            )
+        return False
 
     def execute(self, context: Context) -> Any:
         """Airflow runs this method on the worker and defers using the trigger."""
@@ -464,7 +473,7 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
         self.deferrable = deferrable
         self.poll_interval = poll_interval
 
-    def poke(self, context: Context) -> bool:
+    def poke(self, context: Context) -> PokeReturnValue | bool:
         if self.fail_on_terminal_state:
             job = self.hook.get_job(
                 job_id=self.job_id,
@@ -481,8 +490,16 @@ class DataflowJobAutoScalingEventsSensor(BaseSensorOperator):
             project_id=self.project_id,
             location=self.location,
         )
-
-        return result if self.callback is None else self.callback(result)
+        result = result if self.callback is None else self.callback(result)
+        if isinstance(result, PokeReturnValue):
+            return result
+
+        if bool(result):
+            return PokeReturnValue(
+                is_done=True,
+                xcom_value=result,
+            )
+        return False
 
     def execute(self, context: Context) -> Any:
         """Airflow runs this method on the worker and defers using the trigger."""

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

@@ -23,11 +23,11 @@ from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.dataform import DataformHook
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class DataformWorkflowInvocationStateSensor(BaseSensorOperator):

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

@@ -23,12 +23,12 @@ from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
@@ -110,7 +110,7 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
                 pipeline_id=self.pipeline_id,
                 namespace=self.namespace,
             )
-            pipeline_status = pipeline_workflow["status"]
+            pipeline_status = pipeline_workflow.get("status")
         except AirflowNotFoundException:
             message = "Specified Pipeline ID was not found."
             raise AirflowException(message)
@@ -127,4 +127,4 @@ class CloudDataFusionPipelineStateSensor(BaseSensorOperator):
         self.log.debug(
             "Current status of the pipeline workflow for %s: %s.", self.pipeline_id, pipeline_status
         )
-        return pipeline_status in self.expected_statuses
+        return pipeline_status is not None and pipeline_status in self.expected_statuses

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

@@ -23,19 +23,21 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
+from google.api_core.exceptions import GoogleAPICallError
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.dataplex_v1.types import DataScanJob
+
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.dataplex import (
     AirflowDataQualityScanException,
     AirflowDataQualityScanResultTimeoutException,
     DataplexHook,
 )
-from airflow.sensors.base import BaseSensorOperator
-from google.api_core.exceptions import GoogleAPICallError
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.dataplex_v1.types import DataScanJob
 
 
 class TaskState:

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

@@ -22,11 +22,11 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.dataprep import GoogleDataprepHook, JobGroupStatuses
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class DataprepJobGroupIsFinishedSensor(BaseSensorOperator):

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

@@ -23,15 +23,16 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.exceptions import ServerError
+from google.cloud.dataproc_v1.types import Batch, JobStatus
+
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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 google.api_core.exceptions import ServerError
-from google.cloud.dataproc_v1.types import Batch, JobStatus
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class DataprocJobSensor(BaseSensorOperator):
@@ -99,17 +100,17 @@ class DataprocJobSensor(BaseSensorOperator):
         if state == JobStatus.State.ERROR:
             message = f"Job failed:\n{job}"
             raise AirflowException(message)
-        elif state in {
+        if state in {
             JobStatus.State.CANCELLED,
             JobStatus.State.CANCEL_PENDING,
             JobStatus.State.CANCEL_STARTED,
         }:
             message = f"Job was cancelled:\n{job}"
             raise AirflowException(message)
-        elif state == JobStatus.State.DONE:
+        if state == JobStatus.State.DONE:
             self.log.debug("Job %s completed successfully.", self.dataproc_job_id)
             return True
-        elif state == JobStatus.State.ATTEMPT_FAILURE:
+        if state == JobStatus.State.ATTEMPT_FAILURE:
             self.log.debug("Job %s attempt has failed.", self.dataproc_job_id)
 
         self.log.info("Waiting for job %s to complete.", self.dataproc_job_id)
@@ -178,13 +179,13 @@ class DataprocBatchSensor(BaseSensorOperator):
         if state == Batch.State.FAILED:
             message = "Batch failed"
             raise AirflowException(message)
-        elif state in {
+        if state in {
             Batch.State.CANCELLED,
             Batch.State.CANCELLING,
         }:
             message = "Batch was cancelled."
             raise AirflowException(message)
-        elif state == Batch.State.SUCCEEDED:
+        if state == Batch.State.SUCCEEDED:
             self.log.debug("Batch %s completed successfully.", self.batch_id)
             return True
 

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

@@ -21,14 +21,15 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.operation import Operation
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class MetastoreHivePartitionSensor(BaseSensorOperator):
     """
@@ -111,7 +112,7 @@ class MetastoreHivePartitionSensor(BaseSensorOperator):
 
         # Extract actual query results
         result_base_uri = result_manifest_uri.rsplit("/", 1)[0]
-        results = (f"{result_base_uri}//{filename}" for filename in manifest.get("filenames", []))
+        results = (f"{result_base_uri}/{filename}" for filename in manifest.get("filenames", []))
         found_partitions = sum(
             len(
                 parse_json_from_gcs(

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

@@ -21,12 +21,15 @@ 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
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator, poke_mode_only
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.triggers.gcs import (
     GCSBlobTrigger,
@@ -34,13 +37,12 @@ from airflow.providers.google.cloud.triggers.gcs import (
     GCSPrefixBlobTrigger,
     GCSUploadSessionTrigger,
 )
-from airflow.sensors.base import BaseSensorOperator, poke_mode_only
-from google.cloud.storage.retry import DEFAULT_RETRY
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class GCSObjectExistenceSensor(BaseSensorOperator):
     """
@@ -304,23 +306,22 @@ class GCSObjectsWithPrefixExistenceSensor(BaseSensorOperator):
         if not self.deferrable:
             super().execute(context)
             return self._matches
+        if not self.poke(context=context):
+            self.defer(
+                timeout=timedelta(seconds=self.timeout),
+                trigger=GCSPrefixBlobTrigger(
+                    bucket=self.bucket,
+                    prefix=self.prefix,
+                    poke_interval=self.poke_interval,
+                    google_cloud_conn_id=self.google_cloud_conn_id,
+                    hook_params={
+                        "impersonation_chain": self.impersonation_chain,
+                    },
+                ),
+                method_name="execute_complete",
+            )
         else:
-            if not self.poke(context=context):
-                self.defer(
-                    timeout=timedelta(seconds=self.timeout),
-                    trigger=GCSPrefixBlobTrigger(
-                        bucket=self.bucket,
-                        prefix=self.prefix,
-                        poke_interval=self.poke_interval,
-                        google_cloud_conn_id=self.google_cloud_conn_id,
-                        hook_params={
-                            "impersonation_chain": self.impersonation_chain,
-                        },
-                    ),
-                    method_name="execute_complete",
-                )
-            else:
-                return self._matches
+            return self._matches
 
     def execute_complete(self, context: dict[str, Any], event: dict[str, str | list[str]]) -> str | list[str]:
         """Return immediately and rely on trigger to throw a success event. Callback for the trigger."""

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

@@ -22,11 +22,11 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.looker import JobStatus, LookerHook
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class LookerCheckPdtBuildSensor(BaseSensorOperator):
@@ -65,13 +65,13 @@ class LookerCheckPdtBuildSensor(BaseSensorOperator):
             msg = status_dict["message"]
             message = f'PDT materialization job failed. Job id: {self.materialization_id}. Message:\n"{msg}"'
             raise AirflowException(message)
-        elif status == JobStatus.CANCELLED.value:
+        if status == JobStatus.CANCELLED.value:
             message = f"PDT materialization job was cancelled. Job id: {self.materialization_id}."
             raise AirflowException(message)
-        elif status == JobStatus.UNKNOWN.value:
+        if status == JobStatus.UNKNOWN.value:
             message = f"PDT materialization job has unknown status. Job id: {self.materialization_id}."
             raise AirflowException(message)
-        elif status == JobStatus.DONE.value:
+        if status == JobStatus.DONE.value:
             self.log.debug(
                 "PDT materialization job completed successfully. Job id: %s.", self.materialization_id
             )

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

@@ -19,20 +19,21 @@
 
 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
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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 google.cloud import pubsub_v1
-from google.cloud.pubsub_v1.types import ReceivedMessage
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class PubSubMessageTransformException(AirflowException):
@@ -167,20 +168,19 @@ class PubSubPullSensor(BaseSensorOperator):
         if not self.deferrable:
             super().execute(context)
             return self._return_value
-        else:
-            self.defer(
-                timeout=timedelta(seconds=self.timeout),
-                trigger=PubsubPullTrigger(
-                    project_id=self.project_id,
-                    subscription=self.subscription,
-                    max_messages=self.max_messages,
-                    ack_messages=self.ack_messages,
-                    poke_interval=self.poke_interval,
-                    gcp_conn_id=self.gcp_conn_id,
-                    impersonation_chain=self.impersonation_chain,
-                ),
-                method_name="execute_complete",
-            )
+        self.defer(
+            timeout=timedelta(seconds=self.timeout),
+            trigger=PubsubPullTrigger(
+                project_id=self.project_id,
+                subscription=self.subscription,
+                max_messages=self.max_messages,
+                ack_messages=self.ack_messages,
+                poke_interval=self.poke_interval,
+                gcp_conn_id=self.gcp_conn_id,
+                impersonation_chain=self.impersonation_chain,
+            ),
+            method_name="execute_complete",
+        )
 
     def execute_complete(self, context: Context, event: dict[str, str | list[str]]) -> Any:
         """If messages_callback is provided, execute it; otherwise, return immediately with trigger event message."""

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

@@ -22,12 +22,12 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class TaskQueueEmptySensor(BaseSensorOperator):

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

@@ -24,11 +24,11 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import FeatureStoreHook
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class FeatureViewSyncSensor(BaseSensorOperator):

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

@@ -19,17 +19,19 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.workflows.executions_v1beta import Execution
+
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 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 google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.workflows.executions_v1beta import Execution
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class WorkflowExecutionSensor(BaseSensorOperator):
     """

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

@@ -35,7 +35,7 @@ except ModuleNotFoundError as e:
     raise AirflowOptionalProviderFeatureException(e)
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class ADLSToGCSOperator(ADLSListOperator):

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
@@ -32,7 +32,7 @@ except ModuleNotFoundError as e:
     raise AirflowOptionalProviderFeatureException(e)
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class AzureBlobStorageToGCSOperator(BaseOperator):

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
@@ -34,7 +34,7 @@ except ModuleNotFoundError as e:
     raise AirflowOptionalProviderFeatureException(e)
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class AzureFileShareToGCSOperator(BaseOperator):

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

@@ -22,12 +22,13 @@ 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
+    from airflow.providers.common.compat.sdk import Context
 
 
 class BigQueryToBigQueryOperator(BaseOperator):
@@ -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"],