apache-airflow-providers-google 16.1.0__py3-none-any.whl → 17.0.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/vertex_ai/experiment_service.py

@@ -0,0 +1,435 @@
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from google.api_core import exceptions
+
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.vertex_ai.experiment_service import (
+    ExperimentHook,
+    ExperimentRunHook,
+)
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class CreateExperimentOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to create experiment.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_description: Optional. Description of the evaluation experiment.
+    :param experiment_tensorboard: Optional. The Vertex TensorBoard instance to use as a backing
+        TensorBoard for the provided experiment. If no TensorBoard is provided, a default TensorBoard
+        instance is created and used by this experiment.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        experiment_description: str = "",
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        experiment_tensorboard: str | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.experiment_description = experiment_description
+        self.experiment_tensorboard = experiment_tensorboard
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.create_experiment(
+                project_id=self.project_id,
+                location=self.location,
+                experiment_name=self.experiment_name,
+                experiment_description=self.experiment_description,
+                experiment_tensorboard=self.experiment_tensorboard,
+            )
+        except exceptions.AlreadyExists:
+            raise AirflowException(f"Experiment with name {self.experiment_name} already exist")
+
+        self.log.info("Created experiment: %s", self.experiment_name)
+
+
+class DeleteExperimentOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to delete experiment.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param delete_backing_tensorboard_runs: Optional. If True will also delete the Vertex AI TensorBoard
+            runs associated with the experiment runs under this experiment that we used to store time series
+            metrics.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        delete_backing_tensorboard_runs: bool = False,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.delete_backing_tensorboard_runs = delete_backing_tensorboard_runs
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.delete_experiment(
+                project_id=self.project_id,
+                location=self.location,
+                experiment_name=self.experiment_name,
+                delete_backing_tensorboard_runs=self.delete_backing_tensorboard_runs,
+            )
+        except exceptions.NotFound:
+            raise AirflowException(f"Experiment with name {self.experiment_name} not found")
+
+        self.log.info("Deleted experiment: %s", self.experiment_name)
+
+
+class CreateExperimentRunOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to create experiment run.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_run_name: Required. The specific run name or ID for this experiment.
+    :param experiment_run_tensorboard: Optional. A backing TensorBoard resource to enable and store time series
+        metrics logged to this experiment run using log_time_series_metrics.
+    :param run_after_creation: Optional. If True experiment run will be created with state running.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+        "experiment_run_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        experiment_run_name: str,
+        experiment_run_tensorboard: str | None = None,
+        run_after_creation: bool = False,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.experiment_run_name = experiment_run_name
+        self.experiment_run_tensorboard = experiment_run_tensorboard
+        self.run_after_creation = run_after_creation
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentRunHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.create_experiment_run(
+                project_id=self.project_id,
+                location=self.location,
+                experiment_name=self.experiment_name,
+                experiment_run_name=self.experiment_run_name,
+                experiment_run_tensorboard=self.experiment_run_tensorboard,
+                run_after_creation=self.run_after_creation,
+            )
+        except exceptions.AlreadyExists:
+            raise AirflowException(f"Experiment Run with name {self.experiment_run_name} already exist")
+
+        self.log.info("Created experiment run: %s", self.experiment_run_name)
+
+
+class ListExperimentRunsOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to list experiment runs in experiment.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> list[str]:
+        self.hook = ExperimentRunHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            experiment_runs = self.hook.list_experiment_runs(
+                project_id=self.project_id, experiment_name=self.experiment_name, location=self.location
+            )
+        except exceptions.NotFound:
+            raise AirflowException("Experiment %s not found", self.experiment_name)
+
+        return [er.name for er in experiment_runs]
+
+
+class UpdateExperimentRunStateOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to update state of the experiment run.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_run_name: Required. The specific run name or ID for this experiment.
+    :param new_state: Required. The specific state of experiment run.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+        "experiment_run_name",
+        "new_state",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        experiment_run_name: str,
+        new_state: int,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.experiment_run_name = experiment_run_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+        self.new_state = new_state
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentRunHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.update_experiment_run_state(
+                project_id=self.project_id,
+                experiment_name=self.experiment_name,
+                experiment_run_name=self.experiment_run_name,
+                new_state=self.new_state,
+                location=self.location,
+            )
+            self.log.info("New state of the %s is: %s", self.experiment_run_name, self.new_state)
+        except exceptions.NotFound:
+            raise AirflowException("Experiment or experiment run not found")
+
+
+class DeleteExperimentRunOperator(GoogleCloudBaseOperator):
+    """
+    Use the Vertex AI SDK to delete experiment run.
+
+    :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+    :param location: Required. The ID of the Google Cloud location that the service belongs to.
+    :param experiment_name: Required. The name of the evaluation experiment.
+    :param experiment_run_name: Required. The specific run name or ID for this experiment.
+    :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
+    :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).
+    """
+
+    template_fields = (
+        "location",
+        "project_id",
+        "impersonation_chain",
+        "experiment_name",
+        "experiment_run_name",
+    )
+
+    def __init__(
+        self,
+        *,
+        project_id: str,
+        location: str,
+        experiment_name: str,
+        experiment_run_name: str,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.project_id = project_id
+        self.location = location
+        self.experiment_name = experiment_name
+        self.experiment_run_name = experiment_run_name
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        self.hook = ExperimentRunHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+        try:
+            self.hook.delete_experiment_run(
+                project_id=self.project_id,
+                location=self.location,
+                experiment_name=self.experiment_name,
+                experiment_run_name=self.experiment_run_name,
+            )
+        except exceptions.NotFound:
+            raise AirflowException(f"Experiment Run with name {self.experiment_run_name} not found")
+
+        self.log.info("Deleted experiment run: %s", self.experiment_run_name)

airflow/providers/google/cloud/operators/vertex_ai/generative_model.py

@@ -24,12 +24,13 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from google.api_core import exceptions
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import (
     ExperimentRunHook,
     GenerativeModelHook,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.deprecated import deprecated
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -587,6 +588,11 @@ class GenerateFromCachedContentOperator(GoogleCloudBaseOperator):
         return cached_content_text
 
 
+@deprecated(
+    planned_removal_date="January 3, 2026",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.experiment_service.DeleteExperimentRunOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class DeleteExperimentRunOperator(GoogleCloudBaseOperator):
     """
     Use the Rapid Evaluation API to evaluate a model.

airflow/providers/google/cloud/operators/vertex_ai/ray.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 from google.api_core.exceptions import NotFound
 from google.cloud.aiplatform.vertex_ray.util import resources
@@ -93,8 +93,10 @@ class CreateRayClusterOperator(RayBaseOperator):
     :param location: Required. The ID of the Google Cloud region that the service belongs to.
     :param head_node_type: The head node resource. Resources.node_count must be 1. If not set, default
         value of Resources() class will be used.
-    :param python_version: Python version for the ray cluster.
-    :param ray_version: Ray version for the ray cluster. Default is 2.33.0.
+    :param python_version: Required. Python version for the ray cluster.
+    :param ray_version: Required. Ray version for the ray cluster.
+        Currently only 3 version are available: 2.9.3, 2.33, 2.42. For more information please refer to
+        https://github.com/googleapis/python-aiplatform/blob/main/setup.py#L101
     :param network: Virtual private cloud (VPC) network. For Ray Client, VPC peering is required to
         connect to the Ray Cluster managed in the Vertex API service. For Ray Job API, VPC network is not
         required because Ray Cluster connection can be accessed through dashboard address.
@@ -136,9 +138,9 @@ class CreateRayClusterOperator(RayBaseOperator):
 
     def __init__(
         self,
+        python_version: str,
+        ray_version: Literal["2.9.3", "2.33", "2.42"],
         head_node_type: resources.Resources = resources.Resources(),
-        python_version: str = "3.10",
-        ray_version: str = "2.33",
         network: str | None = None,
         service_account: str | None = None,
         cluster_name: str | None = None,

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

@@ -693,7 +693,7 @@ class CloudVisionUpdateProductOperator(GoogleCloudBaseOperator):
             location=self.location,
             product_id=self.product_id,
             project_id=self.project_id,
-            update_mask=self.update_mask,  # type: ignore
+            update_mask=self.update_mask,
             retry=self.retry,
             timeout=self.timeout,
             metadata=self.metadata,

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

@@ -37,7 +37,7 @@ from airflow.providers.google.cloud.triggers.dataflow import (
     DataflowJobStatusTrigger,
 )
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from airflow.providers.google.version_compat import BaseSensorOperator
+from airflow.providers.google.version_compat import BaseSensorOperator, PokeReturnValue
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -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/datafusion.py

@@ -115,7 +115,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)
@@ -132,4 +132,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/transfers/gcs_to_bigquery.py

@@ -429,7 +429,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):
@@ -581,7 +580,7 @@ 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,
             location=self.location,

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

@@ -20,13 +20,15 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-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
 
+# 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/oracle_to_gcs.py

@@ -46,15 +46,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):