diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py index cefc5a24b46d..93120f3d075e 100644 --- a/airflow/providers/google/cloud/hooks/bigtable.py +++ b/airflow/providers/google/cloud/hooks/bigtable.py @@ -22,12 +22,11 @@ import warnings from typing import Sequence -from google.cloud.bigtable import Client +from google.cloud.bigtable import Client, enums from google.cloud.bigtable.cluster import Cluster from google.cloud.bigtable.column_family import ColumnFamily, GarbageCollectionRule from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.table import ClusterState, Table -from google.cloud.bigtable_admin_v2 import enums from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook @@ -56,9 +55,9 @@ def __init__( delegate_to=delegate_to, impersonation_chain=impersonation_chain, ) - self._client = None + self._client: Client | None = None - def _get_client(self, project_id: str): + def _get_client(self, project_id: str) -> Client: if not self._client: self._client = Client( project=project_id, @@ -69,7 +68,7 @@ def _get_client(self, project_id: str): return self._client @GoogleBaseHook.fallback_to_default_project_id - def get_instance(self, instance_id: str, project_id: str) -> Instance: + def get_instance(self, instance_id: str, project_id: str) -> Instance | None: """ Retrieves and returns the specified Cloud Bigtable instance if it exists. Otherwise, returns None. @@ -113,10 +112,10 @@ def create_instance( project_id: str, replica_clusters: list[dict[str, str]] | None = None, instance_display_name: str | None = None, - instance_type: enums.Instance.Type = enums.Instance.Type.TYPE_UNSPECIFIED, + instance_type: enums.Instance.Type = enums.Instance.Type.UNSPECIFIED, # type: ignore[assignment] instance_labels: dict | None = None, cluster_nodes: int | None = None, - cluster_storage_type: enums.StorageType = enums.StorageType.STORAGE_TYPE_UNSPECIFIED, + cluster_storage_type: enums.StorageType = enums.StorageType.UNSPECIFIED, # type: ignore[assignment] timeout: float | None = None, ) -> Instance: """ @@ -142,9 +141,6 @@ def create_instance( :param timeout: (optional) timeout (in seconds) for instance creation. If None is not specified, Operator will wait indefinitely. """ - cluster_storage_type = enums.StorageType(cluster_storage_type) - instance_type = enums.Instance.Type(instance_type) - instance = Instance( instance_id, self._get_client(project_id=project_id), @@ -200,8 +196,6 @@ def update_instance( :param timeout: (optional) timeout (in seconds) for instance update. If None is not specified, Operator will wait indefinitely. """ - instance_type = enums.Instance.Type(instance_type) - instance = Instance( instance_id=instance_id, client=self._get_client(project_id=project_id), @@ -253,7 +247,10 @@ def delete_table(self, instance_id: str, table_id: str, project_id: str) -> None BigTable exists. If set to None or missing, the default project_id from the Google Cloud connection is used. """ - table = self.get_instance(instance_id=instance_id, project_id=project_id).table(table_id=table_id) + instance = self.get_instance(instance_id=instance_id, project_id=project_id) + if instance is None: + raise RuntimeError("Instance %s did not exist; unable to delete table %s" % instance_id, table_id) + table = instance.table(table_id=table_id) table.delete() @staticmethod diff --git a/airflow/providers/google/cloud/operators/bigtable.py b/airflow/providers/google/cloud/operators/bigtable.py index e2e54166d546..50c86c56a8f9 100644 --- a/airflow/providers/google/cloud/operators/bigtable.py +++ b/airflow/providers/google/cloud/operators/bigtable.py @@ -22,8 +22,8 @@ from typing import TYPE_CHECKING, Iterable, Sequence import google.api_core.exceptions +from google.cloud.bigtable import enums from google.cloud.bigtable.column_family import GarbageCollectionRule -from google.cloud.bigtable_admin_v2 import enums from airflow.exceptions import AirflowException from airflow.models import BaseOperator diff --git a/airflow/providers/google/cloud/sensors/bigtable.py b/airflow/providers/google/cloud/sensors/bigtable.py index 3dd689f8c608..c69d4099cad1 100644 --- a/airflow/providers/google/cloud/sensors/bigtable.py +++ b/airflow/providers/google/cloud/sensors/bigtable.py @@ -21,8 +21,8 @@ from typing import TYPE_CHECKING, Sequence import google.api_core.exceptions +from google.cloud.bigtable import enums from google.cloud.bigtable.table import ClusterState -from google.cloud.bigtable_admin_v2 import enums from airflow.providers.google.cloud.hooks.bigtable import BigtableHook from airflow.providers.google.cloud.links.bigtable import BigtableTablesLink @@ -103,7 +103,7 @@ def poke(self, context: Context) -> bool: ) return False - ready_state = ClusterState(enums.Table.ClusterState.ReplicationState.READY) + ready_state = ClusterState(enums.Table.ReplicationState.READY) is_table_replicated = True for cluster_id in cluster_states.keys(): diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index 3326a9cb8d31..1feed66374bd 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -83,7 +83,7 @@ dependencies: - google-cloud-aiplatform>=1.7.1,<2.0.0 - google-cloud-automl>=2.1.0 - google-cloud-bigquery-datatransfer>=3.0.0 - - google-cloud-bigtable>=1.0.0,<2.0.0 + - google-cloud-bigtable>=2.0.0,<3.0.0 - google-cloud-build>=3.0.0 - google-cloud-compute>=0.1.0,<2.0.0 - google-cloud-container>=2.2.0,<3.0.0 @@ -127,9 +127,6 @@ dependencies: # A transient dependency of google-cloud-bigquery-datatransfer, but we # further constrain it since older versions are buggy. - proto-plus>=1.19.6 - # Google bigtable client require protobuf <= 3.20.0. We can remove the limitation - # when this limitation is removed - - protobuf<=3.20.0 integrations: - integration-name: Google Analytics360 diff --git a/docs/apache-airflow-providers-google/index.rst b/docs/apache-airflow-providers-google/index.rst index d19d58506d44..e5d5b1c45911 100644 --- a/docs/apache-airflow-providers-google/index.rst +++ b/docs/apache-airflow-providers-google/index.rst @@ -112,7 +112,7 @@ PIP package Version required ``google-cloud-aiplatform`` ``>=1.7.1,<2.0.0`` ``google-cloud-automl`` ``>=2.1.0`` ``google-cloud-bigquery-datatransfer`` ``>=3.0.0`` -``google-cloud-bigtable`` ``>=1.0.0,<2.0.0`` +``google-cloud-bigtable`` ``>=2.0.0,<3.0.0`` ``google-cloud-build`` ``>=3.0.0`` ``google-cloud-compute`` ``>=0.1.0,<2.0.0`` ``google-cloud-container`` ``>=2.2.0,<3.0.0`` diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 6a914f9ca514..575efe0def3e 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -336,7 +336,7 @@ "google-cloud-aiplatform>=1.7.1,<2.0.0", "google-cloud-automl>=2.1.0", "google-cloud-bigquery-datatransfer>=3.0.0", - "google-cloud-bigtable>=1.0.0,<2.0.0", + "google-cloud-bigtable>=2.0.0,<3.0.0", "google-cloud-build>=3.0.0", "google-cloud-compute>=0.1.0,<2.0.0", "google-cloud-container>=2.2.0,<3.0.0", @@ -373,7 +373,6 @@ "pandas-gbq", "pandas>=0.17.1", "proto-plus>=1.19.6", - "protobuf<=3.20.0", "sqlalchemy-bigquery>=1.2.1" ], "cross-providers-deps": [ diff --git a/tests/providers/google/cloud/hooks/test_bigtable.py b/tests/providers/google/cloud/hooks/test_bigtable.py index 93e77a36ac65..f3f2049f797d 100644 --- a/tests/providers/google/cloud/hooks/test_bigtable.py +++ b/tests/providers/google/cloud/hooks/test_bigtable.py @@ -21,9 +21,8 @@ from unittest.mock import PropertyMock import google -from google.cloud.bigtable import Client +from google.cloud.bigtable import Client, enums from google.cloud.bigtable.instance import Instance -from google.cloud.bigtable_admin_v2 import enums from airflow.providers.google.cloud.hooks.bigtable import BigtableHook from airflow.providers.google.common.consts import CLIENT_INFO diff --git a/tests/providers/google/cloud/operators/test_bigtable.py b/tests/providers/google/cloud/operators/test_bigtable.py index 68e5134c8b88..ffb7051289a3 100644 --- a/tests/providers/google/cloud/operators/test_bigtable.py +++ b/tests/providers/google/cloud/operators/test_bigtable.py @@ -21,9 +21,9 @@ import google.api_core.exceptions import pytest +from google.cloud.bigtable import enums from google.cloud.bigtable.column_family import MaxVersionsGCRule from google.cloud.bigtable.instance import Instance -from google.cloud.bigtable_admin_v2 import enums from airflow.exceptions import AirflowException from airflow.providers.google.cloud.operators.bigtable import (