apache-airflow-providers-google 12.0.0rc1__py3-none-any.whl → 13.0.0__py3-none-any.whl

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

@@ -23,11 +23,6 @@ import uuid
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.workflows.executions_v1beta import Execution
-from google.cloud.workflows_v1beta import Workflow
-
 from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook
 from airflow.providers.google.cloud.links.workflows import (
     WorkflowsExecutionLink,
@@ -36,13 +31,16 @@ from airflow.providers.google.cloud.links.workflows import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.workflows.executions_v1beta import Execution
+from google.cloud.workflows_v1beta import Workflow
 
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
-
 from airflow.utils.hashlib_wrapper import md5
 
 

airflow/providers/google/cloud/secrets/secret_manager.py

@@ -21,8 +21,6 @@ from __future__ import annotations
 import logging
 from collections.abc import Sequence
 
-from google.auth.exceptions import DefaultCredentialsError
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud._internal_client.secret_manager_client import _SecretManagerClient
 from airflow.providers.google.cloud.utils.credentials_provider import (
@@ -32,6 +30,7 @@ from airflow.providers.google.cloud.utils.credentials_provider import (
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.secrets import BaseSecretsBackend
 from airflow.utils.log.logging_mixin import LoggingMixin
+from google.auth.exceptions import DefaultCredentialsError
 
 log = logging.getLogger(__name__)
 

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

@@ -22,18 +22,16 @@ 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.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 google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.bigquery_datatransfer_v1 import TransferState
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 class BigQueryDataTransferServiceTransferRunSensor(BaseSensorOperator):

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

@@ -23,14 +23,13 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 import google.api_core.exceptions
-from google.cloud.bigtable import enums
-from google.cloud.bigtable.table import ClusterState
-
 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 google.cloud.bigtable import enums
+from google.cloud.bigtable.table import ClusterState
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,10 +22,10 @@ from __future__ import annotations
 import json
 from collections.abc import Iterable, Sequence
 from datetime import datetime, timedelta
+from functools import cached_property
 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
@@ -34,6 +34,7 @@ from airflow.providers.google.cloud.triggers.cloud_composer import CloudComposer
 from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
 from airflow.sensors.base import BaseSensorOperator
 from airflow.utils.state import TaskInstanceState
+from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -135,19 +136,20 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
 
     def _pull_dag_runs(self) -> list[dict]:
         """Pull the list of dag runs."""
-        hook = CloudComposerHook(
-            gcp_conn_id=self.gcp_conn_id,
-            impersonation_chain=self.impersonation_chain,
+        cmd_parameters = (
+            ["-d", self.composer_dag_id, "-o", "json"]
+            if self._composer_airflow_version < 3
+            else [self.composer_dag_id, "-o", "json"]
         )
-        dag_runs_cmd = hook.execute_airflow_command(
+        dag_runs_cmd = self.hook.execute_airflow_command(
             project_id=self.project_id,
             region=self.region,
             environment_id=self.environment_id,
             command="dags",
             subcommand="list-runs",
-            parameters=["-d", self.composer_dag_id, "-o", "json"],
+            parameters=cmd_parameters,
         )
-        cmd_result = hook.wait_command_execution_result(
+        cmd_result = self.hook.wait_command_execution_result(
             project_id=self.project_id,
             region=self.region,
             environment_id=self.environment_id,
@@ -165,13 +167,27 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
         for dag_run in dag_runs:
             if (
                 start_date.timestamp()
-                < parser.parse(dag_run["logical_date"]).timestamp()
+                < parser.parse(
+                    dag_run["execution_date" if self._composer_airflow_version < 3 else "logical_date"]
+                ).timestamp()
                 < end_date.timestamp()
             ) and dag_run["state"] not in self.allowed_states:
                 return False
         return True
 
+    def _get_composer_airflow_version(self) -> int:
+        """Return Composer Airflow version."""
+        environment_obj = self.hook.get_environment(
+            project_id=self.project_id,
+            region=self.region,
+            environment_id=self.environment_id,
+        )
+        environment_config = Environment.to_dict(environment_obj)
+        image_version = environment_config["config"]["software_config"]["image_version"]
+        return int(image_version.split("airflow-")[1].split(".")[0])
+
     def execute(self, context: Context) -> None:
+        self._composer_airflow_version = self._get_composer_airflow_version()
         if self.deferrable:
             start_date, end_date = self._get_logical_dates(context)
             self.defer(
@@ -186,6 +202,7 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
                     gcp_conn_id=self.gcp_conn_id,
                     impersonation_chain=self.impersonation_chain,
                     poll_interval=self.poll_interval,
+                    composer_airflow_version=self._composer_airflow_version,
                 ),
                 method_name=GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME,
             )
@@ -195,3 +212,10 @@ class CloudComposerDAGRunSensor(BaseSensorOperator):
         if event and event["status"] == "error":
             raise AirflowException(event["message"])
         self.log.info("DAG %s has executed successfully.", self.composer_dag_id)
+
+    @cached_property
+    def hook(self) -> CloudComposerHook:
+        return CloudComposerHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )

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

@@ -20,8 +20,10 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     COUNTERS,
     METADATA,
@@ -29,6 +31,9 @@ from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import
     CloudDataTransferServiceHook,
 )
 from airflow.providers.google.cloud.links.cloud_storage_transfer import CloudStorageTransferJobLink
+from airflow.providers.google.cloud.triggers.cloud_storage_transfer_service import (
+    CloudStorageTransferServiceCheckJobStatusTrigger,
+)
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.sensors.base import BaseSensorOperator
 
@@ -60,6 +65,7 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
         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 deferrable: Run sensor in deferrable mode
     """
 
     # [START gcp_transfer_job_sensor_template_fields]
@@ -78,6 +84,7 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
         project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -88,6 +95,7 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
         self.project_id = project_id
         self.gcp_cloud_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
+        self.deferrable = deferrable
 
     def poke(self, context: Context) -> bool:
         hook = CloudDataTransferServiceHook(
@@ -117,3 +125,33 @@ class CloudDataTransferServiceJobStatusSensor(BaseSensorOperator):
             )
 
         return check
+
+    def execute(self, context: Context) -> None:
+        """Run on the worker and defer using the triggers if deferrable is set to True."""
+        if not self.deferrable:
+            super().execute(context)
+        elif not self.poke(context=context):
+            self.defer(
+                timeout=self.execution_timeout,
+                trigger=CloudStorageTransferServiceCheckJobStatusTrigger(
+                    job_name=self.job_name,
+                    expected_statuses=self.expected_statuses,
+                    project_id=self.project_id,
+                    poke_interval=self.poke_interval,
+                    gcp_conn_id=self.gcp_cloud_conn_id,
+                    impersonation_chain=self.impersonation_chain,
+                ),
+                method_name="execute_complete",
+            )
+
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
+        """
+        Act as a callback for when the trigger fires.
+
+        This returns immediately. It relies on trigger to throw an exception,
+        otherwise it assumes execution was successful.
+        """
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+
+        self.xcom_push(key="sensed_operations", value=event["operations"], context=context)

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

@@ -23,13 +23,8 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context 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 google.api_core.retry import Retry
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.dataplex import (
@@ -38,6 +33,9 @@ from airflow.providers.google.cloud.hooks.dataplex import (
     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/dataproc.py

@@ -23,13 +23,12 @@ 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.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

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

@@ -26,9 +26,8 @@ from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
 from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
-    from google.api_core.operation import Operation
-
     from airflow.utils.context import Context
+    from google.api_core.operation import Operation
 
 
 class MetastoreHivePartitionSensor(BaseSensorOperator):

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

@@ -25,8 +25,6 @@ from collections.abc import Sequence
 from datetime import datetime, timedelta
 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
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
@@ -37,11 +35,11 @@ from airflow.providers.google.cloud.triggers.gcs import (
     GCSUploadSessionTrigger,
 )
 from airflow.sensors.base import BaseSensorOperator, poke_mode_only
+from google.cloud.storage.retry import DEFAULT_RETRY
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 class GCSObjectExistenceSensor(BaseSensorOperator):

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

@@ -23,14 +23,13 @@ from collections.abc import Sequence
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Callable
 
-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.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

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

@@ -19,18 +19,16 @@ 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.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 google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 class WorkflowExecutionSensor(BaseSensorOperator):

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

@@ -22,9 +22,6 @@ from __future__ import annotations
 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
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator
@@ -33,11 +30,12 @@ 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.utils.helpers import merge_dicts
+from google.api_core.exceptions import Conflict
+from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 class BigQueryToGCSOperator(BaseOperator):
@@ -222,6 +220,8 @@ class BigQueryToGCSOperator(BaseOperator):
 
         try:
             self.log.info("Executing: %s", configuration)
+            if not self.job_id:
+                raise AirflowException("Job ID is None, expected a valid job ID.")
             job: BigQueryJob | UnknownJob = self._submit_job(
                 hook=hook, job_id=self.job_id, configuration=configuration
             )

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

@@ -23,6 +23,15 @@ import json
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+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.utils.helpers import merge_dicts
 from google.api_core.exceptions import BadRequest, Conflict
 from google.cloud.bigquery import (
     DEFAULT_RETRY,
@@ -36,20 +45,9 @@ from google.cloud.bigquery import (
 )
 from google.cloud.bigquery.table import EncryptionConfiguration, Table, TableReference
 
-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.utils.helpers import merge_dicts
-
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 ALLOWED_FORMATS = [
     "CSV",

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

@@ -304,7 +304,7 @@ class GCSToGCSOperator(BaseOperator):
                 dest_object.replace(self.destination_object, prefix, 1) for dest_object in destination_objects
             ]
 
-        objects = set(objects) - set(existing_objects)
+        objects = list(set(objects) - set(existing_objects))
         if objects:
             self.log.info("%s files are going to be synced: %s.", len(objects), objects)
         else:

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

@@ -21,6 +21,7 @@ from __future__ import annotations
 
 import os
 from collections.abc import Sequence
+from functools import cached_property
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
@@ -129,14 +130,16 @@ class GCSToSFTPOperator(BaseOperator):
         self.impersonation_chain = impersonation_chain
         self.sftp_dirs = None
 
+    @cached_property
+    def sftp_hook(self):
+        return SFTPHook(self.sftp_conn_id)
+
     def execute(self, context: Context):
         gcs_hook = GCSHook(
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
 
-        sftp_hook = SFTPHook(self.sftp_conn_id)
-
         if WILDCARD in self.source_object:
             total_wildcards = self.source_object.count(WILDCARD)
             if total_wildcards > 1:
@@ -155,12 +158,12 @@ class GCSToSFTPOperator(BaseOperator):
 
             for source_object in objects:
                 destination_path = self._resolve_destination_path(source_object, prefix=prefix_dirname)
-                self._copy_single_object(gcs_hook, sftp_hook, source_object, destination_path)
+                self._copy_single_object(gcs_hook, self.sftp_hook, source_object, destination_path)
 
             self.log.info("Done. Uploaded '%d' files to %s", len(objects), self.destination_path)
         else:
             destination_path = self._resolve_destination_path(self.source_object)
-            self._copy_single_object(gcs_hook, sftp_hook, self.source_object, destination_path)
+            self._copy_single_object(gcs_hook, self.sftp_hook, self.source_object, destination_path)
             self.log.info("Done. Uploaded '%s' file to %s", self.source_object, destination_path)
 
     def _resolve_destination_path(self, source_object: str, prefix: str | None = None) -> str:
@@ -200,3 +203,32 @@ class GCSToSFTPOperator(BaseOperator):
         if self.move_object:
             self.log.info("Executing delete of gs://%s/%s", self.source_bucket, source_object)
             gcs_hook.delete(self.source_bucket, source_object)
+
+    def get_openlineage_facets_on_start(self):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        source_name = extract_ds_name_from_gcs_path(f"{self.source_object}")
+        dest_name = f"{self.destination_path}"
+        if self.keep_directory_structure:
+            dest_name = os.path.join(dest_name, source_name if source_name != "/" else "")
+        elif WILDCARD not in self.source_object:
+            dest_name = os.path.join(dest_name, os.path.basename(self.source_object))
+
+        dest_name = dest_name.rstrip("/") if dest_name != "/" else "/"
+
+        return OperatorLineage(
+            inputs=[
+                Dataset(
+                    namespace=f"gs://{self.source_bucket}",
+                    name=source_name,
+                )
+            ],
+            outputs=[
+                Dataset(
+                    namespace=f"file://{self.sftp_hook.remote_host}:{self.sftp_hook.port}",
+                    name=dest_name,
+                )
+            ],
+        )

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

@@ -22,10 +22,15 @@ from __future__ import annotations
 import datetime
 import decimal
 from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.microsoft.mssql.hooks.mssql import MsSqlHook
 
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class MSSQLToGCSOperator(BaseSQLToGCSOperator):
     """
@@ -75,14 +80,17 @@ class MSSQLToGCSOperator(BaseSQLToGCSOperator):
         self.mssql_conn_id = mssql_conn_id
         self.bit_fields = bit_fields or []
 
+    @cached_property
+    def db_hook(self) -> MsSqlHook:
+        return MsSqlHook(mssql_conn_id=self.mssql_conn_id)
+
     def query(self):
         """
         Query MSSQL and returns a cursor of results.
 
         :return: mssql cursor
         """
-        mssql = MsSqlHook(mssql_conn_id=self.mssql_conn_id)
-        conn = mssql.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor()
         cursor.execute(self.sql)
         return cursor
@@ -109,3 +117,20 @@ class MSSQLToGCSOperator(BaseSQLToGCSOperator):
         if isinstance(value, (datetime.date, datetime.time)):
             return value.isoformat()
         return value
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        sql_parsing_result = get_openlineage_facets_with_sql(
+            hook=self.db_hook,
+            sql=self.sql,
+            conn_id=self.mssql_conn_id,
+            database=None,
+        )
+        gcs_output_datasets = self._get_openlineage_output_datasets()
+        if sql_parsing_result:
+            sql_parsing_result.outputs = gcs_output_datasets
+            return sql_parsing_result
+        return OperatorLineage(outputs=gcs_output_datasets, job_facets={"sql": SQLJobFacet(self.sql)})

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

@@ -22,6 +22,8 @@ from __future__ import annotations
 import base64
 from datetime import date, datetime, time, timedelta
 from decimal import Decimal
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 try:
     from MySQLdb.constants import FIELD_TYPE
@@ -37,6 +39,9 @@ except ImportError:
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.mysql.hooks.mysql import MySqlHook
 
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class MySQLToGCSOperator(BaseSQLToGCSOperator):
     """
@@ -77,10 +82,13 @@ class MySQLToGCSOperator(BaseSQLToGCSOperator):
         self.mysql_conn_id = mysql_conn_id
         self.ensure_utc = ensure_utc
 
+    @cached_property
+    def db_hook(self) -> MySqlHook:
+        return MySqlHook(mysql_conn_id=self.mysql_conn_id)
+
     def query(self):
         """Query mysql and returns a cursor to the results."""
-        mysql = MySqlHook(mysql_conn_id=self.mysql_conn_id)
-        conn = mysql.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor()
         if self.ensure_utc:
             # Ensure TIMESTAMP results are in UTC
@@ -140,3 +148,20 @@ class MySQLToGCSOperator(BaseSQLToGCSOperator):
             else:
                 value = base64.standard_b64encode(value).decode("ascii")
         return value
+
+    def get_openlineage_facets_on_start(self) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        sql_parsing_result = get_openlineage_facets_with_sql(
+            hook=self.db_hook,
+            sql=self.sql,
+            conn_id=self.mysql_conn_id,
+            database=None,
+        )
+        gcs_output_datasets = self._get_openlineage_output_datasets()
+        if sql_parsing_result:
+            sql_parsing_result.outputs = gcs_output_datasets
+            return sql_parsing_result
+        return OperatorLineage(outputs=gcs_output_datasets, job_facets={"sql": SQLJobFacet(self.sql)})