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

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

@@ -23,15 +23,6 @@ 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,
@@ -45,10 +36,21 @@ 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.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:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 ALLOWED_FORMATS = [
     "CSV",
     "NEWLINE_DELIMITED_JSON",
@@ -56,6 +58,7 @@ ALLOWED_FORMATS = [
     "GOOGLE_SHEETS",
     "DATASTORE_BACKUP",
     "PARQUET",
+    "ORC",
 ]
 
 
@@ -141,6 +144,9 @@ class GCSToBigQueryOperator(BaseOperator):
         partition by field, type and  expiration as per API specifications.
         Note that 'field' is not available in concurrency with
         dataset.table$partition.
+        Ignored if 'range_partitioning' is set.
+    :param range_partitioning: configure optional range partitioning fields i.e.
+        partition by field and integer interval as per API specifications.
     :param cluster_fields: Request that the result of this load be stored sorted
         by one or more columns. BigQuery supports clustering for both partitioned and
         non-partitioned tables. The order of columns given determines the sort order.
@@ -216,6 +222,7 @@ class GCSToBigQueryOperator(BaseOperator):
         src_fmt_configs=None,
         external_table=False,
         time_partitioning=None,
+        range_partitioning=None,
         cluster_fields=None,
         autodetect=True,
         encryption_configuration=None,
@@ -243,6 +250,10 @@ class GCSToBigQueryOperator(BaseOperator):
             src_fmt_configs = {}
         if time_partitioning is None:
             time_partitioning = {}
+        if range_partitioning is None:
+            range_partitioning = {}
+        if range_partitioning and time_partitioning:
+            raise ValueError("Only one of time_partitioning or range_partitioning can be set.")
         self.bucket = bucket
         self.source_objects = source_objects
         self.schema_object = schema_object
@@ -260,8 +271,7 @@ class GCSToBigQueryOperator(BaseOperator):
                 f"{source_format} is not a valid source format. "
                 f"Please use one of the following types: {ALLOWED_FORMATS}."
             )
-        else:
-            self.source_format = source_format.upper()
+        self.source_format = source_format.upper()
         self.compression = compression
         self.create_disposition = create_disposition
         self.skip_leading_rows = skip_leading_rows
@@ -281,6 +291,7 @@ class GCSToBigQueryOperator(BaseOperator):
         self.schema_update_options = schema_update_options
         self.src_fmt_configs = src_fmt_configs
         self.time_partitioning = time_partitioning
+        self.range_partitioning = range_partitioning
         self.cluster_fields = cluster_fields
         self.autodetect = autodetect
         self.encryption_configuration = encryption_configuration
@@ -335,8 +346,9 @@ class GCSToBigQueryOperator(BaseOperator):
             job_id=self.job_id,
             dag_id=self.dag_id,
             task_id=self.task_id,
-            logical_date=context["logical_date"],
+            logical_date=None,
             configuration=self.configuration,
+            run_after=hook.get_run_after_or_logical_date(context),
             force_rerun=self.force_rerun,
         )
 
@@ -371,7 +383,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"],
@@ -404,14 +415,13 @@ class GCSToBigQueryOperator(BaseOperator):
                         f"want to force rerun it consider setting `force_rerun=True`."
                         f"Or, if you want to reattach in this scenario add {job.state} to `reattach_states`"
                     )
-                else:
-                    # Job already reached state DONE
-                    if job.state == "DONE":
-                        raise AirflowException("Job is already in state DONE. Can not reattach to this job.")
+                # Job already reached state DONE
+                if job.state == "DONE":
+                    raise AirflowException("Job is already in state DONE. Can not reattach to this job.")
 
-                    # We are reattaching to a job
-                    self.log.info("Reattaching to existing Job in state %s", job.state)
-                    self._handle_job_error(job)
+                # We are reattaching to a job
+                self.log.info("Reattaching to existing Job in state %s", job.state)
+                self._handle_job_error(job)
 
             job_types = {
                 LoadJob._JOB_TYPE: ["sourceTable", "destinationTable"],
@@ -429,7 +439,6 @@ class GCSToBigQueryOperator(BaseOperator):
                                 table = job_configuration[job_type][table_prop]
                                 persist_kwargs = {
                                     "context": context,
-                                    "task_instance": self,
                                     "table_id": table,
                                 }
                                 if not isinstance(table, str):
@@ -484,8 +493,7 @@ class GCSToBigQueryOperator(BaseOperator):
         if self.max_id_key:
             self.log.info("Selecting the MAX value from BigQuery column %r...", self.max_id_key)
             select_command = (
-                f"SELECT MAX({self.max_id_key}) AS max_value "
-                f"FROM {self.destination_project_dataset_table}"
+                f"SELECT MAX({self.max_id_key}) AS max_value FROM {self.destination_project_dataset_table}"
             )
             self.configuration = {
                 "query": {
@@ -505,8 +513,7 @@ class GCSToBigQueryOperator(BaseOperator):
                         f"Could not determine MAX value in column {self.max_id_key} "
                         f"since the default value of 'string_field_n' was set by BQ"
                     )
-                else:
-                    raise AirflowException(e.message)
+                raise AirflowException(e.message)
             if rows:
                 for row in rows:
                     max_id = row[0] if row[0] else 0
@@ -548,6 +555,7 @@ class GCSToBigQueryOperator(BaseOperator):
                 "quote",
                 "encoding",
                 "preserveAsciiControlCharacters",
+                "columnNameCharacterMap",
             ],
             "googleSheetsOptions": ["skipLeadingRows"],
         }
@@ -582,9 +590,11 @@ class GCSToBigQueryOperator(BaseOperator):
         table_obj_api_repr = table.to_api_repr()
 
         self.log.info("Creating external table: %s", self.destination_project_dataset_table)
-        self.hook.create_empty_table(
+        self.hook.create_table(
             table_resource=table_obj_api_repr,
             project_id=self.project_id or self.hook.project_id,
+            dataset_id=table.dataset_id,
+            table_id=table.table_id,
             location=self.location,
             exists_ok=True,
         )
@@ -629,6 +639,8 @@ class GCSToBigQueryOperator(BaseOperator):
         )
         if self.time_partitioning:
             self.configuration["load"].update({"timePartitioning": self.time_partitioning})
+        if self.range_partitioning:
+            self.configuration["load"].update({"rangePartitioning": self.range_partitioning})
 
         if self.cluster_fields:
             self.configuration["load"].update({"clustering": {"fields": self.cluster_fields}})
@@ -637,17 +649,10 @@ class GCSToBigQueryOperator(BaseOperator):
             self.configuration["load"]["schema"] = {"fields": self.schema_fields}
 
         if self.schema_update_options:
-            if self.write_disposition not in ["WRITE_APPEND", "WRITE_TRUNCATE"]:
-                raise ValueError(
-                    "schema_update_options is only "
-                    "allowed if write_disposition is "
-                    "'WRITE_APPEND' or 'WRITE_TRUNCATE'."
-                )
-            else:
-                # To provide backward compatibility
-                self.schema_update_options = list(self.schema_update_options or [])
-                self.log.info("Adding experimental 'schemaUpdateOptions': %s", self.schema_update_options)
-                self.configuration["load"]["schemaUpdateOptions"] = self.schema_update_options
+            # To provide backward compatibility
+            self.schema_update_options = list(self.schema_update_options or [])
+            self.log.info("Adding experimental 'schemaUpdateOptions': %s", self.schema_update_options)
+            self.configuration["load"]["schemaUpdateOptions"] = self.schema_update_options
 
         if self.max_bad_records:
             self.configuration["load"]["maxBadRecords"] = self.max_bad_records
@@ -674,11 +679,13 @@ class GCSToBigQueryOperator(BaseOperator):
                 "quote",
                 "encoding",
                 "preserveAsciiControlCharacters",
+                "columnNameCharacterMap",
             ],
             "DATASTORE_BACKUP": ["projectionFields"],
             "NEWLINE_DELIMITED_JSON": ["autodetect", "ignoreUnknownValues"],
             "PARQUET": ["autodetect", "ignoreUnknownValues"],
             "AVRO": ["useAvroLogicalTypes"],
+            "ORC": ["autodetect"],
         }
 
         valid_configs = src_fmt_to_configs_mapping[self.source_format]

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

@@ -24,13 +24,13 @@ 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 = "*"
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class GCSToGCSOperator(BaseOperator):

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

@@ -20,12 +20,14 @@ 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
+    from airflow.providers.common.compat.sdk import Context
+
+# MAX XCOM Size is 48KB, check discussion: https://github.com/apache/airflow/pull/1618#discussion_r68249677
+MAX_XCOM_SIZE = 49344
 
 
 class GCSToLocalFilesystemOperator(BaseOperator):

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

@@ -26,14 +26,14 @@ 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 = "*"
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class GCSToSFTPOperator(BaseOperator):
@@ -80,7 +80,9 @@ class GCSToSFTPOperator(BaseOperator):
     :param destination_path: The sftp remote path. This is the specified directory path for
         uploading to the SFTP server.
     :param keep_directory_structure: (Optional) When set to False the path of the file
-         on the bucket is recreated within path passed in destination_path.
+        on the bucket is recreated within path passed in destination_path.
+    :param create_intermediate_dirs: (Optional) When set to True the intermediate directories
+        in the specified file path will be created.
     :param move_object: When move object is True, the object is moved instead
         of copied to the new location. This is the equivalent of a mv command
         as opposed to a cp command.
@@ -112,6 +114,7 @@ class GCSToSFTPOperator(BaseOperator):
         source_object: str,
         destination_path: str,
         keep_directory_structure: bool = True,
+        create_intermediate_dirs: bool = True,
         move_object: bool = False,
         gcp_conn_id: str = "google_cloud_default",
         sftp_conn_id: str = "ssh_default",
@@ -124,6 +127,7 @@ class GCSToSFTPOperator(BaseOperator):
         self.source_object = source_object
         self.destination_path = destination_path
         self.keep_directory_structure = keep_directory_structure
+        self.create_intermediate_dirs = create_intermediate_dirs
         self.move_object = move_object
         self.gcp_conn_id = gcp_conn_id
         self.sftp_conn_id = sftp_conn_id
@@ -190,7 +194,9 @@ class GCSToSFTPOperator(BaseOperator):
         )
 
         dir_path = os.path.dirname(destination_path)
-        sftp_hook.create_directory(dir_path)
+
+        if self.create_intermediate_dirs:
+            sftp_hook.create_directory(dir_path)
 
         with NamedTemporaryFile("w") as tmp:
             gcs_hook.download(

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

@@ -19,12 +19,12 @@ 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
+    from airflow.providers.common.compat.sdk import Context
 
 
 class GoogleDriveToGCSOperator(BaseOperator):
@@ -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

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

@@ -19,11 +19,11 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 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
+    from airflow.providers.common.compat.sdk import Context
 
 
 class GoogleDriveToLocalOperator(BaseOperator):

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

@@ -0,0 +1,193 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""This module contains operator to move data from HTTP endpoint to GCS."""
+
+from __future__ import annotations
+
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.version_compat import BaseOperator
+from airflow.providers.http.hooks.http import HttpHook
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from requests.auth import AuthBase
+
+    from airflow.providers.common.compat.sdk import Context
+
+
+class HttpToGCSOperator(BaseOperator):
+    """
+    Calls an endpoint on an HTTP system to execute an action and store the result in GCS.
+
+    :param http_conn_id: The :ref:`http connection<howto/connection:http>` to run
+        the operator against
+    :param endpoint: The relative part of the full url. (templated)
+    :param method: The HTTP method to use, default = "POST"
+    :param data: The data to pass. POST-data in POST/PUT and params
+        in the URL for a GET request. (templated)
+    :param headers: The HTTP headers to be added to the GET request
+    :param response_check: A check against the 'requests' response object.
+        The callable takes the response object as the first positional argument
+        and optionally any number of keyword arguments available in the context dictionary.
+        It should return True for 'pass' and False otherwise.
+    :param response_filter: A function allowing you to manipulate the response
+        text. e.g response_filter=lambda response: json.loads(response.text).
+        The callable takes the response object as the first positional argument
+        and optionally any number of keyword arguments available in the context dictionary.
+    :param extra_options: Extra options for the 'requests' library, see the
+        'requests' documentation (options to modify timeout, ssl, etc.)
+    :param log_response: Log the response (default: False)
+    :param auth_type: The auth type for the service
+    :param tcp_keep_alive: Enable TCP Keep Alive for the connection.
+    :param tcp_keep_alive_idle: The TCP Keep Alive Idle parameter (corresponds to ``socket.TCP_KEEPIDLE``).
+    :param tcp_keep_alive_count: The TCP Keep Alive count parameter (corresponds to ``socket.TCP_KEEPCNT``)
+    :param tcp_keep_alive_interval: The TCP Keep Alive interval parameter (corresponds to
+        ``socket.TCP_KEEPINTVL``)
+    :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.
+    :param bucket_name: The bucket to upload to.
+    :param object_name: The object name to set when uploading the file.
+    :param mime_type: The file mime type set when uploading the file.
+    :param gzip: Option to compress local file or file data for upload
+    :param encoding: bytes encoding for file data if provided as string
+    :param chunk_size: Blob chunk size.
+    :param timeout: Request timeout in seconds.
+    :param num_max_attempts: Number of attempts to try to upload the file.
+    :param metadata: The metadata to be uploaded with the file.
+    :param cache_contro: Cache-Control metadata field.
+    :param user_project: The identifier of the Google Cloud project to bill for the request. Required for Requester Pays buckets.
+    """
+
+    template_fields: Sequence[str] = (
+        "http_conn_id",
+        "endpoint",
+        "data",
+        "headers",
+        "gcp_conn_id",
+        "bucket_name",
+        "object_name",
+    )
+    template_fields_renderers = {"headers": "json", "data": "py"}
+    template_ext: Sequence[str] = ()
+    ui_color = "#f4a460"
+
+    def __init__(
+        self,
+        *,
+        endpoint: str | None = None,
+        method: str = "GET",
+        data: Any = None,
+        headers: dict[str, str] | None = None,
+        extra_options: dict[str, Any] | None = None,
+        http_conn_id: str = "http_default",
+        log_response: bool = False,
+        auth_type: type[AuthBase] | None = None,
+        tcp_keep_alive: bool = True,
+        tcp_keep_alive_idle: int = 120,
+        tcp_keep_alive_count: int = 20,
+        tcp_keep_alive_interval: int = 30,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        bucket_name: str,
+        object_name: str,
+        mime_type: str | None = None,
+        gzip: bool = False,
+        encoding: str | None = None,
+        chunk_size: int | None = None,
+        timeout: int | None = None,
+        num_max_attempts: int = 3,
+        metadata: dict | None = None,
+        cache_control: str | None = None,
+        user_project: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.http_conn_id = http_conn_id
+        self.method = method
+        self.endpoint = endpoint
+        self.headers = headers or {}
+        self.data = data or {}
+        self.extra_options = extra_options or {}
+        self.log_response = log_response
+        self.auth_type = auth_type
+        self.tcp_keep_alive = tcp_keep_alive
+        self.tcp_keep_alive_idle = tcp_keep_alive_idle
+        self.tcp_keep_alive_count = tcp_keep_alive_count
+        self.tcp_keep_alive_interval = tcp_keep_alive_interval
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.bucket_name = bucket_name
+        self.object_name = object_name
+        self.mime_type = mime_type
+        self.gzip = gzip
+        self.encoding = encoding
+        self.chunk_size = chunk_size
+        self.timeout = timeout
+        self.num_max_attempts = num_max_attempts
+        self.metadata = metadata
+        self.cache_control = cache_control
+        self.user_project = user_project
+
+    @cached_property
+    def http_hook(self) -> HttpHook:
+        """Create and return an HttpHook."""
+        return HttpHook(
+            self.method,
+            http_conn_id=self.http_conn_id,
+            auth_type=self.auth_type,
+            tcp_keep_alive=self.tcp_keep_alive,
+            tcp_keep_alive_idle=self.tcp_keep_alive_idle,
+            tcp_keep_alive_count=self.tcp_keep_alive_count,
+            tcp_keep_alive_interval=self.tcp_keep_alive_interval,
+        )
+
+    @cached_property
+    def gcs_hook(self) -> GCSHook:
+        """Create and return an GCSHook."""
+        return GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
+
+    def execute(self, context: Context):
+        self.log.info("Calling HTTP method")
+        response = self.http_hook.run(
+            endpoint=self.endpoint, data=self.data, headers=self.headers, extra_options=self.extra_options
+        )
+
+        self.log.info("Uploading to GCS")
+        self.gcs_hook.upload(
+            data=response.content,
+            bucket_name=self.bucket_name,
+            object_name=self.object_name,
+            mime_type=self.mime_type,
+            gzip=self.gzip,
+            encoding=self.encoding or response.encoding,
+            chunk_size=self.chunk_size,
+            timeout=self.timeout,
+            num_max_attempts=self.num_max_attempts,
+            metadata=self.metadata,
+            cache_control=self.cache_control,
+            user_project=self.user_project,
+        )

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

@@ -24,11 +24,11 @@ from collections.abc import Sequence
 from glob import glob
 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
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class LocalFilesystemToGCSOperator(BaseOperator):

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

@@ -67,7 +67,7 @@ class MSSQLToGCSOperator(BaseSQLToGCSOperator):
 
     ui_color = "#e0a98c"
 
-    type_map = {2: "BOOLEAN", 3: "INTEGER", 4: "TIMESTAMP", 5: "NUMERIC"}
+    type_map = {2: "BOOL", 3: "INTEGER", 4: "TIMESTAMP", 5: "NUMERIC"}
 
     def __init__(
         self,

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

@@ -21,12 +21,17 @@ import base64
 import calendar
 from datetime import date, datetime, timedelta
 from decimal import Decimal
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 import oracledb
 
 from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator
 from airflow.providers.oracle.hooks.oracle import OracleHook
 
+if TYPE_CHECKING:
+    from airflow.providers.openlineage.extractors import OperatorLineage
+
 
 class OracleToGCSOperator(BaseSQLToGCSOperator):
     """
@@ -46,15 +51,15 @@ class OracleToGCSOperator(BaseSQLToGCSOperator):
     ui_color = "#a0e08c"
 
     type_map = {
-        oracledb.DB_TYPE_BINARY_DOUBLE: "DECIMAL",  # type: ignore
-        oracledb.DB_TYPE_BINARY_FLOAT: "DECIMAL",  # type: ignore
-        oracledb.DB_TYPE_BINARY_INTEGER: "INTEGER",  # type: ignore
-        oracledb.DB_TYPE_BOOLEAN: "BOOLEAN",  # type: ignore
-        oracledb.DB_TYPE_DATE: "TIMESTAMP",  # type: ignore
-        oracledb.DB_TYPE_NUMBER: "NUMERIC",  # type: ignore
-        oracledb.DB_TYPE_TIMESTAMP: "TIMESTAMP",  # type: ignore
-        oracledb.DB_TYPE_TIMESTAMP_LTZ: "TIMESTAMP",  # type: ignore
-        oracledb.DB_TYPE_TIMESTAMP_TZ: "TIMESTAMP",  # type: ignore
+        oracledb.DB_TYPE_BINARY_DOUBLE: "DECIMAL",
+        oracledb.DB_TYPE_BINARY_FLOAT: "DECIMAL",
+        oracledb.DB_TYPE_BINARY_INTEGER: "INTEGER",
+        oracledb.DB_TYPE_BOOLEAN: "BOOLEAN",
+        oracledb.DB_TYPE_DATE: "TIMESTAMP",
+        oracledb.DB_TYPE_NUMBER: "NUMERIC",
+        oracledb.DB_TYPE_TIMESTAMP: "TIMESTAMP",
+        oracledb.DB_TYPE_TIMESTAMP_LTZ: "TIMESTAMP",
+        oracledb.DB_TYPE_TIMESTAMP_TZ: "TIMESTAMP",
     }
 
     def __init__(self, *, oracle_conn_id="oracle_default", ensure_utc=False, **kwargs):
@@ -62,10 +67,13 @@ class OracleToGCSOperator(BaseSQLToGCSOperator):
         self.ensure_utc = ensure_utc
         self.oracle_conn_id = oracle_conn_id
 
+    @cached_property
+    def db_hook(self) -> OracleHook:
+        return OracleHook(oracle_conn_id=self.oracle_conn_id)
+
     def query(self):
         """Query Oracle and returns a cursor to the results."""
-        oracle = OracleHook(oracle_conn_id=self.oracle_conn_id)
-        conn = oracle.get_conn()
+        conn = self.db_hook.get_conn()
         cursor = conn.cursor()
         if self.ensure_utc:
             # Ensure TIMESTAMP results are in UTC
@@ -121,3 +129,20 @@ class OracleToGCSOperator(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.oracle_conn_id,
+            database=self.db_hook.service_name or self.db_hook.sid,
+        )
+        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)})