apache-airflow-providers-google 10.26.0rc1__py3-none-any.whl → 11.0.0rc1__py3-none-any.whl

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

@@ -20,21 +20,18 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterable, Sequence
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Iterable, Sequence
+from typing import TYPE_CHECKING
 
 from dateutil import parser
 from google.cloud.orchestration.airflow.service_v1.types import ExecuteAirflowCommandResponse
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+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,
-    CloudComposerExecutionTrigger,
-)
+from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposerDAGRunTrigger
 from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.sensors.base import BaseSensorOperator
 from airflow.utils.state import TaskInstanceState
 
@@ -42,89 +39,6 @@ if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="CloudComposerCreateEnvironmentOperator, CloudComposerDeleteEnvironmentOperator, "
-    "CloudComposerUpdateEnvironmentOperator",
-    instructions="Please use CloudComposerCreateEnvironmentOperator, CloudComposerDeleteEnvironmentOperator "
-    "or CloudComposerUpdateEnvironmentOperator in deferrable or non-deferrable mode, "
-    "since since every operator gives user a possibility to wait (asynchronously or synchronously) "
-    "until the Operation is finished.",
-    category=AirflowProviderDeprecationWarning,
-)
-class CloudComposerEnvironmentSensor(BaseSensorOperator):
-    """
-    Check the status of the Cloud Composer Environment task.
-
-    This Sensor is deprecated. You can achieve the same functionality by using Cloud Composer Operators
-    CloudComposerCreateEnvironmentOperator, CloudComposerDeleteEnvironmentOperator and
-    CloudComposerUpdateEnvironmentOperator in  deferrable or non-deferrable mode, since every operator
-    gives user a possibility to wait (asynchronously or synchronously) until Operation will be finished.
-
-    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
-    :param region: Required. The ID of the Google Cloud region that the service belongs to.
-    :param operation_name: The name of the operation resource
-    :param gcp_conn_id: The connection ID to use when fetching connection info.
-    :param impersonation_chain: Optional service account to impersonate using short-term
-        credentials, or chained list of accounts required to get the access_token
-        of the last account in the list, which will be impersonated in the request.
-        If set as a string, the account must grant the originating account
-        the Service Account Token Creator IAM role.
-        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 pooling_period_seconds: Optional: Control the rate of the poll for the result of deferrable run.
-    """
-
-    def __init__(
-        self,
-        *,
-        project_id: str,
-        region: str,
-        operation_name: str,
-        gcp_conn_id: str = "google_cloud_default",
-        impersonation_chain: str | Sequence[str] | None = None,
-        pooling_period_seconds: int = 30,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.project_id = project_id
-        self.region = region
-        self.operation_name = operation_name
-        self.pooling_period_seconds = pooling_period_seconds
-        self.gcp_conn_id = gcp_conn_id
-        self.impersonation_chain = impersonation_chain
-
-    def execute(self, context: Context) -> None:
-        """Airflow runs this method on the worker and defers using the trigger."""
-        self.defer(
-            trigger=CloudComposerExecutionTrigger(
-                project_id=self.project_id,
-                region=self.region,
-                operation_name=self.operation_name,
-                gcp_conn_id=self.gcp_conn_id,
-                impersonation_chain=self.impersonation_chain,
-                pooling_period_seconds=self.pooling_period_seconds,
-            ),
-            method_name="execute_complete",
-        )
-
-    def execute_complete(self, context: dict[str, Any], event: dict[str, str] | None = None) -> str:
-        """
-        Act as a callback for when the trigger fires - returns immediately.
-
-        Relies on trigger to throw an exception, otherwise it assumes execution was successful.
-        """
-        if event:
-            if event.get("operation_done"):
-                return event["operation_done"]
-
-            raise AirflowException(event["message"])
-
-        message = "No event received in trigger callback"
-        raise AirflowException(message)
-
-
 class CloudComposerDAGRunSensor(BaseSensorOperator):
     """
     Check if a DAG run has completed.
@@ -189,7 +103,7 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
         self.deferrable = deferrable
         self.poll_interval = poll_interval
 
-    def _get_execution_dates(self, context) -> tuple[datetime, datetime]:
+    def _get_logical_dates(self, context) -> tuple[datetime, datetime]:
         if isinstance(self.execution_range, timedelta):
             if self.execution_range < timedelta(0):
                 return context["logical_date"], context["logical_date"] - self.execution_range
@@ -203,7 +117,7 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
             return context["logical_date"] - timedelta(1), context["logical_date"]
 
     def poke(self, context: Context) -> bool:
-        start_date, end_date = self._get_execution_dates(context)
+        start_date, end_date = self._get_logical_dates(context)
 
         if datetime.now(end_date.tzinfo) < end_date:
             return False
@@ -251,7 +165,7 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
         for dag_run in dag_runs:
             if (
                 start_date.timestamp()
-                < parser.parse(dag_run["execution_date"]).timestamp()
+                < parser.parse(dag_run["logical_date"]).timestamp()
                 < end_date.timestamp()
             ) and dag_run["state"] not in self.allowed_states:
                 return False
@@ -259,7 +173,7 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
 
     def execute(self, context: Context) -> None:
         if self.deferrable:
-            start_date, end_date = self._get_execution_dates(context)
+            start_date, end_date = self._get_logical_dates(context)
             self.defer(
                 trigger=CloudComposerDAGRunTrigger(
                     project_id=self.project_id,

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     COUNTERS,

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

@@ -19,8 +19,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable, Sequence
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataform import DataformHook

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable, Sequence
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowNotFoundException
 from airflow.providers.google.cloud.hooks.datafusion import DataFusionHook

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

@@ -19,7 +19,8 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+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

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import time
-from typing import TYPE_CHECKING, Sequence
+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

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook

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

@@ -21,13 +21,14 @@ from __future__ import annotations
 
 import os
 import textwrap
+from collections.abc import Sequence
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable
 
 from google.cloud.storage.retry import DEFAULT_RETRY
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.triggers.gcs import (
     GCSBlobTrigger,
@@ -35,7 +36,6 @@ from airflow.providers.google.cloud.triggers.gcs import (
     GCSPrefixBlobTrigger,
     GCSUploadSessionTrigger,
 )
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.sensors.base import BaseSensorOperator, poke_mode_only
 
 if TYPE_CHECKING:
@@ -142,38 +142,6 @@ class GCSObjectExistenceSensor(BaseSensorOperator):
         return True
 
 
-@deprecated(
-    planned_removal_date="November 01, 2024",
-    use_instead="GCSObjectExistenceSensor",
-    instructions="Please use GCSObjectExistenceSensor and set deferrable attribute to True.",
-    category=AirflowProviderDeprecationWarning,
-)
-class GCSObjectExistenceAsyncSensor(GCSObjectExistenceSensor):
-    """
-    Checks for the existence of a file in Google Cloud Storage.
-
-    This class is deprecated and will be removed in a future release.
-
-    Please use :class:`airflow.providers.google.cloud.sensors.gcs.GCSObjectExistenceSensor`
-    and set *deferrable* attribute to *True* instead.
-
-    :param bucket: The Google Cloud Storage bucket where the object is.
-    :param object: The name of the object to check in the Google cloud storage bucket.
-    :param google_cloud_conn_id: The connection ID to use when connecting to Google Cloud Storage.
-    :param impersonation_chain: Optional service account to impersonate using short-term
-        credentials, or chained list of accounts required to get the access_token
-        of the last account in the list, which will be impersonated in the request.
-        If set as a string, the account must grant the originating account
-        the Service Account Token Creator IAM role.
-        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).
-    """
-
-    def __init__(self, **kwargs: Any) -> None:
-        super().__init__(deferrable=True, **kwargs)
-
-
 def ts_function(context):
     """
     Act as a default callback for the GoogleCloudStorageObjectUpdatedSensor.
@@ -192,7 +160,7 @@ class GCSObjectUpdateSensor(BaseSensorOperator):
     :param object: The name of the object to download in the Google cloud
         storage bucket.
     :param ts_func: Callback for defining the update condition. The default callback
-        returns execution_date + schedule_interval. The callback takes the context
+        returns logical_date + schedule_interval. The callback takes the context
         as parameter.
     :param google_cloud_conn_id: The connection ID to use when
         connecting to Google Cloud Storage.

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

@@ -19,8 +19,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable
 
 from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1.types import ReceivedMessage

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+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

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+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

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

@@ -20,8 +20,9 @@
 from __future__ import annotations
 
 import os
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
 

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

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import tempfile
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook

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

@@ -18,8 +18,9 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
@@ -110,6 +111,7 @@ class BigQueryToBigQueryOperator(BaseOperator):
         self.location = location
         self.impersonation_chain = impersonation_chain
         self.hook: BigQueryHook | None = None
+        self._job_conf: dict = {}
 
     def _prepare_job_configuration(self):
         self.source_project_dataset_tables = (
@@ -154,39 +156,94 @@ class BigQueryToBigQueryOperator(BaseOperator):
 
         return configuration
 
-    def _submit_job(
-        self,
-        hook: BigQueryHook,
-        configuration: dict,
-    ) -> str:
-        job = hook.insert_job(configuration=configuration, project_id=hook.project_id)
-        return job.job_id
-
     def execute(self, context: Context) -> None:
         self.log.info(
             "Executing copy of %s into: %s",
             self.source_project_dataset_tables,
             self.destination_project_dataset_table,
         )
-        hook = BigQueryHook(
+        self.hook = BigQueryHook(
             gcp_conn_id=self.gcp_conn_id,
             location=self.location,
             impersonation_chain=self.impersonation_chain,
         )
-        self.hook = hook
 
-        if not hook.project_id:
+        if not self.hook.project_id:
             raise ValueError("The project_id should be set")
 
         configuration = self._prepare_job_configuration()
-        job_id = self._submit_job(hook=hook, configuration=configuration)
+        self._job_conf = self.hook.insert_job(
+            configuration=configuration, project_id=self.hook.project_id
+        ).to_api_repr()
 
-        job = hook.get_job(job_id=job_id, location=self.location).to_api_repr()
-        conf = job["configuration"]["copy"]["destinationTable"]
+        dest_table_info = self._job_conf["configuration"]["copy"]["destinationTable"]
         BigQueryTableLink.persist(
             context=context,
             task_instance=self,
-            dataset_id=conf["datasetId"],
-            project_id=conf["projectId"],
-            table_id=conf["tableId"],
+            dataset_id=dest_table_info["datasetId"],
+            project_id=dest_table_info["projectId"],
+            table_id=dest_table_info["tableId"],
+        )
+
+    def get_openlineage_facets_on_complete(self, task_instance):
+        """Implement on_complete as we will include final BQ job id."""
+        from airflow.providers.common.compat.openlineage.facet import (
+            Dataset,
+            ExternalQueryRunFacet,
+        )
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            get_facets_from_bq_table,
+            get_identity_column_lineage_facet,
         )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        if not self.hook:
+            self.hook = BigQueryHook(
+                gcp_conn_id=self.gcp_conn_id,
+                location=self.location,
+                impersonation_chain=self.impersonation_chain,
+            )
+
+        if not self._job_conf:
+            self.log.debug("OpenLineage could not find BQ job configuration.")
+            return OperatorLineage()
+
+        bq_job_id = self._job_conf["jobReference"]["jobId"]
+        source_tables_info = self._job_conf["configuration"]["copy"]["sourceTables"]
+        dest_table_info = self._job_conf["configuration"]["copy"]["destinationTable"]
+
+        run_facets = {
+            "externalQuery": ExternalQueryRunFacet(externalQueryId=bq_job_id, source="bigquery"),
+        }
+
+        input_datasets = []
+        for in_table_info in source_tables_info:
+            table_id = ".".join(
+                (in_table_info["projectId"], in_table_info["datasetId"], in_table_info["tableId"])
+            )
+            table_object = self.hook.get_client().get_table(table_id)
+            input_datasets.append(
+                Dataset(
+                    namespace=BIGQUERY_NAMESPACE, name=table_id, facets=get_facets_from_bq_table(table_object)
+                )
+            )
+
+        out_table_id = ".".join(
+            (dest_table_info["projectId"], dest_table_info["datasetId"], dest_table_info["tableId"])
+        )
+        out_table_object = self.hook.get_client().get_table(out_table_id)
+        output_dataset_facets = {
+            **get_facets_from_bq_table(out_table_object),
+            **get_identity_column_lineage_facet(
+                dest_field_names=[field.name for field in out_table_object.schema],
+                input_datasets=input_datasets,
+            ),
+        }
+        output_dataset = Dataset(
+            namespace=BIGQUERY_NAMESPACE,
+            name=out_table_id,
+            facets=output_dataset_facets,
+        )
+
+        return OperatorLineage(inputs=input_datasets, outputs=[output_dataset], run_facets=run_facets)

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import Conflict
 from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
@@ -294,6 +295,7 @@ class BigQueryToGCSOperator(BaseOperator):
         from pathlib import Path
 
         from airflow.providers.common.compat.openlineage.facet import (
+            BaseFacet,
             Dataset,
             ExternalQueryRunFacet,
             Identifier,
@@ -322,12 +324,12 @@ class BigQueryToGCSOperator(BaseOperator):
             facets=get_facets_from_bq_table(table_object),
         )
 
-        output_dataset_facets = {
-            "schema": input_dataset.facets["schema"],
-            "columnLineage": get_identity_column_lineage_facet(
-                field_names=[field.name for field in table_object.schema], input_datasets=[input_dataset]
-            ),
-        }
+        output_dataset_facets: dict[str, BaseFacet] = get_identity_column_lineage_facet(
+            dest_field_names=[field.name for field in table_object.schema], input_datasets=[input_dataset]
+        )
+        if "schema" in input_dataset.facets:
+            output_dataset_facets["schema"] = input_dataset.facets["schema"]
+
         output_datasets = []
         for uri in sorted(self.destination_cloud_storage_uris):
             bucket, blob = _parse_gcs_url(uri)

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.links.bigquery import BigQueryTableLink

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

@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Sequence
+from collections.abc import Sequence
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.transfers.bigquery_to_sql import BigQueryToSqlBaseOperator

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook

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

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.models import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook

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

@@ -21,10 +21,11 @@ from __future__ import annotations
 
 import json
 from base64 import b64encode
+from collections.abc import Iterable, Sequence
 from datetime import datetime
 from decimal import Decimal
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Any, Iterable, NewType, Sequence
+from typing import TYPE_CHECKING, Any, NewType
 from uuid import UUID
 
 from cassandra.util import Date, OrderedMapSerializedKey, SortedSet, Time

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

@@ -21,8 +21,9 @@ from __future__ import annotations
 
 import csv
 import tempfile
+from collections.abc import Sequence
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator

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

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import BadRequest, Conflict
 from google.cloud.bigquery import (
@@ -784,9 +785,10 @@ class GCSToBigQueryOperator(BaseOperator):
         source_objects = (
             self.source_objects if isinstance(self.source_objects, list) else [self.source_objects]
         )
-        input_dataset_facets = {
-            "schema": output_dataset_facets["schema"],
-        }
+        input_dataset_facets = {}
+        if "schema" in output_dataset_facets:
+            input_dataset_facets["schema"] = output_dataset_facets["schema"]
+
         input_datasets = []
         for blob in sorted(source_objects):
             additional_facets = {}
@@ -811,14 +813,16 @@ class GCSToBigQueryOperator(BaseOperator):
             )
             input_datasets.append(dataset)
 
-        output_dataset_facets["columnLineage"] = get_identity_column_lineage_facet(
-            field_names=[field.name for field in table_object.schema], input_datasets=input_datasets
-        )
-
         output_dataset = Dataset(
             namespace="bigquery",
             name=str(table_object.reference),
-            facets=output_dataset_facets,
+            facets={
+                **output_dataset_facets,
+                **get_identity_column_lineage_facet(
+                    dest_field_names=[field.name for field in table_object.schema],
+                    input_datasets=input_datasets,
+                ),
+            },
         )
 
         run_facets = {}