feat(bigtable): add Bigtable cluster metadata tools

PiperOrigin-RevId: 878648015

Author: google-genai-bot

contributing/samples/bigtable/agent.py

@@ -62,8 +62,8 @@
     credentials_config=credentials_config, bigtable_tool_settings=tool_settings
 )
 
-_BIGTABLE_PROJECT_ID = "google.com:cloud-bigtable-dev"
-_BIGTABLE_INSTANCE_ID = "annenguyen-bus-instance"
+_BIGTABLE_PROJECT_ID = ""
+_BIGTABLE_INSTANCE_ID = ""
 
 
 def search_hotels_by_location(

src/google/adk/tools/bigtable/bigtable_toolset.py

@@ -44,6 +44,8 @@ class BigtableToolset(BaseToolset):
     - bigtable_get_instance_info
     - bigtable_list_tables
     - bigtable_get_table_info
+    - bigtable_list_clusters
+    - bigtable_get_cluster_info
     - bigtable_execute_sql
   """
 
@@ -95,6 +97,8 @@ async def get_tools(
             metadata_tool.get_instance_info,
             metadata_tool.list_tables,
             metadata_tool.get_table_info,
+            metadata_tool.list_clusters,
+            metadata_tool.get_cluster_info,
             query_tool.execute_sql,
         ]
     ]

src/google/adk/tools/bigtable/metadata_tool.py

@@ -14,12 +14,16 @@
 
 from __future__ import annotations
 
+import enum
 import logging
 
 from google.auth.credentials import Credentials
+from google.cloud.bigtable import enums
 
 from . import client
 
+logger = logging.getLogger(f"google_adk.{__name__}")
+
 
 def list_instances(project_id: str, credentials: Credentials) -> dict:
   """List Bigtable instance ids in a Google Cloud project.
@@ -29,7 +33,22 @@ def list_instances(project_id: str, credentials: Credentials) -> dict:
       credentials (Credentials): The credentials to use for the request.
 
   Returns:
-      dict: Dictionary with a list of the Bigtable instance ids present in the project.
+      dict: Dictionary with a list of dictionaries, each representing a Bigtable instance.
+
+      Example:
+        {
+          "status": "SUCCESS",
+          "results": [
+              {
+                  "project_id": "test-project",
+                  "instance_id": "test-instance",
+                  "display_name": "Test Instance",
+                  "state": "READY",
+                  "type": "PRODUCTION",
+                  "labels": {"env": "test"},
+              }
+          ],
+        }
   """
   try:
     bt_client = client.get_bigtable_admin_client(
@@ -41,12 +60,27 @@ def list_instances(project_id: str, credentials: Credentials) -> dict:
           "Failed to list instances from the following locations: %s",
           failed_locations_list,
       )
-    instance_ids = [instance.instance_id for instance in instances_list]
-    return {"status": "SUCCESS", "results": instance_ids}
+    result = [
+        {
+            "project_id": project_id,
+            "instance_id": instance.instance_id,
+            "display_name": instance.display_name,
+            "state": _enum_name_from_value(
+                enums.Instance.State, instance.state, "UNKNOWN_STATE"
+            ),
+            "type": _enum_name_from_value(
+                enums.Instance.Type, instance.type_, "UNKNOWN_TYPE"
+            ),
+            "labels": instance.labels,
+        }
+        for instance in instances_list
+    ]
+    return {"status": "SUCCESS", "results": result}
   except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
     return {
         "status": "ERROR",
-        "error_details": str(ex),
+        "error_details": repr(ex),
     }
 
 
@@ -69,52 +103,85 @@ def get_instance_info(
     )
     instance = bt_client.instance(instance_id)
     instance.reload()
-    instance_info = {
-        "project_id": project_id,
-        "instance_id": instance.instance_id,
-        "display_name": instance.display_name,
-        "state": instance.state,
-        "type": instance.type_,
-        "labels": instance.labels,
+    return {
+        "status": "SUCCESS",
+        "results": {
+            "project_id": project_id,
+            "instance_id": instance.instance_id,
+            "display_name": instance.display_name,
+            "state": _enum_name_from_value(
+                enums.Instance.State, instance.state, "UNKNOWN_STATE"
+            ),
+            "type": _enum_name_from_value(
+                enums.Instance.Type, instance.type_, "UNKNOWN_TYPE"
+            ),
+            "labels": instance.labels,
+        },
     }
-    return {"status": "SUCCESS", "results": instance_info}
   except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
     return {
         "status": "ERROR",
-        "error_details": str(ex),
+        "error_details": repr(ex),
     }
 
 
 def list_tables(
     project_id: str, instance_id: str, credentials: Credentials
 ) -> dict:
-  """List table ids in a Bigtable instance.
+  """List tables and their metadata in a Bigtable instance.
 
   Args:
       project_id (str): The Google Cloud project id containing the instance.
       instance_id (str): The Bigtable instance id.
       credentials (Credentials): The credentials to use for the request.
 
   Returns:
-      dict: Dictionary with a list of the tables ids present in the instance.
+      dict: A dictionary with status and results, where results is a list of
+      table properties.
+
+      Example:
+        {
+          "status": "SUCCESS",
+          "results": [
+              {
+                  "project_id": "test-project",
+                  "instance_id": "test-instance",
+                  "table_id": "test-table",
+                  "table_name": "fake-table-name",
+              }
+          ],
+        }
   """
   try:
     bt_client = client.get_bigtable_admin_client(
         project=project_id, credentials=credentials
     )
     instance = bt_client.instance(instance_id)
     tables = instance.list_tables()
-    table_ids = [table.table_id for table in tables]
-    return {"status": "SUCCESS", "results": table_ids}
+    result = [
+        {
+            "project_id": project_id,
+            "instance_id": instance_id,
+            "table_id": table.table_id,
+            "table_name": table.name,
+        }
+        for table in tables
+    ]
+    return {"status": "SUCCESS", "results": result}
   except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
     return {
         "status": "ERROR",
-        "error_details": str(ex),
+        "error_details": repr(ex),
     }
 
 
 def get_table_info(
-    project_id: str, instance_id: str, table_id: str, credentials: Credentials
+    project_id: str,
+    instance_id: str,
+    table_id: str,
+    credentials: Credentials,
 ) -> dict:
   """Get metadata information about a Bigtable table.
 
@@ -126,6 +193,17 @@ def get_table_info(
 
   Returns:
       dict: Dictionary representing the properties of the table.
+
+      Example:
+        {
+          "status": "SUCCESS",
+          "results": {
+              "project_id": "test-project",
+              "instance_id": "test-instance",
+              "table_id": "test-table",
+              "column_families": ["cf1", "cf2"],
+          },
+        }
   """
   try:
     bt_client = client.get_bigtable_admin_client(
@@ -134,15 +212,170 @@ def get_table_info(
     instance = bt_client.instance(instance_id)
     table = instance.table(table_id)
     column_families = table.list_column_families()
-    table_info = {
-        "project_id": project_id,
-        "instance_id": instance.instance_id,
-        "table_id": table.table_id,
-        "column_families": list(column_families.keys()),
+    return {
+        "status": "SUCCESS",
+        "results": {
+            "project_id": project_id,
+            "instance_id": instance.instance_id,
+            "table_id": table.table_id,
+            "column_families": list(column_families.keys()),
+        },
+    }
+  except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
+    return {
+        "status": "ERROR",
+        "error_details": repr(ex),
+    }
+
+
+def _enum_name_from_value(
+    enum_class: type[enum.Enum], value: int, prefix: str = "UNKNOWN"
+) -> str:
+  for attr_name in dir(enum_class):
+    if not attr_name.startswith("_"):
+      if getattr(enum_class, attr_name) == value:
+        return attr_name
+  return f"{prefix}_{value}"
+
+
+def list_clusters(
+    project_id: str, instance_id: str, credentials: Credentials
+) -> dict:
+  """List clusters and their metadata in a Bigtable instance.
+
+  Args:
+      project_id (str): The Google Cloud project id containing the instance.
+      instance_id (str): The Bigtable instance id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary representing the properties of the cluster.
+
+      Example:
+        {
+          "status": "SUCCESS",
+          "results": [
+              {
+                  "project_id": "test-project",
+                  "instance_id": "test-instance",
+                  "cluster_id": "test-cluster",
+                  "cluster_name": "fake-cluster-name",
+                  "state": "READY",
+                  "serve_nodes": 3,
+                  "default_storage_type": "SSD",
+                  "location_id": "us-central1-a",
+              }
+          ],
+        }
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    instance = bt_client.instance(instance_id)
+    instance.reload()
+    clusters_list, failed_locations = instance.list_clusters()
+    if failed_locations:
+      logging.warning(
+          "Failed to list clusters from the following locations: %s",
+          failed_locations,
+      )
+
+    result = [
+        {
+            "project_id": project_id,
+            "instance_id": instance_id,
+            "cluster_id": cluster.cluster_id,
+            "cluster_name": cluster.name,
+            "state": _enum_name_from_value(
+                enums.Cluster.State, cluster.state, "UNKNOWN_STATE"
+            ),
+            "serve_nodes": cluster.serve_nodes,
+            "default_storage_type": _enum_name_from_value(
+                enums.StorageType,
+                cluster.default_storage_type,
+                "UNKNOWN_STORAGE_TYPE",
+            ),
+            "location_id": cluster.location_id,
+        }
+        for cluster in clusters_list
+    ]
+    return {"status": "SUCCESS", "results": result}
+  except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
+    return {
+        "status": "ERROR",
+        "error_details": repr(ex),
+    }
+
+
+def get_cluster_info(
+    project_id: str,
+    instance_id: str,
+    cluster_id: str,
+    credentials: Credentials,
+) -> dict:
+  """Get detailed metadata information about a Bigtable cluster.
+
+  Args:
+      project_id (str): The Google Cloud project id containing the instance.
+      instance_id (str): The Bigtable instance id containing the cluster.
+      cluster_id (str): The Bigtable cluster id.
+      credentials (Credentials): The credentials to use for the request.
+
+  Returns:
+      dict: Dictionary representing the properties of the cluster.
+
+      Example:
+        {
+          "status": "SUCCESS",
+          "results": {
+              "project_id": "test-project",
+              "instance_id": "test-instance",
+              "cluster_id": "test-cluster",
+              "state": "READY",
+              "serve_nodes": 3,
+              "default_storage_type": "SSD",
+              "location_id": "us-central1-a",
+              "min_serve_nodes": 1,
+              "max_serve_nodes": 10,
+              "cpu_utilization_percent": 80,
+          },
+        }
+  """
+  try:
+    bt_client = client.get_bigtable_admin_client(
+        project=project_id, credentials=credentials
+    )
+    instance = bt_client.instance(instance_id)
+    instance.reload()
+    cluster = instance.cluster(cluster_id)
+    cluster.reload()
+    return {
+        "status": "SUCCESS",
+        "results": {
+            "project_id": project_id,
+            "instance_id": instance_id,
+            "cluster_id": cluster.cluster_id,
+            "state": _enum_name_from_value(
+                enums.Cluster.State, cluster.state, "UNKNOWN_STATE"
+            ),
+            "serve_nodes": cluster.serve_nodes,
+            "default_storage_type": _enum_name_from_value(
+                enums.StorageType,
+                cluster.default_storage_type,
+                "UNKNOWN_STORAGE_TYPE",
+            ),
+            "location_id": cluster.location_id,
+            "min_serve_nodes": cluster.min_serve_nodes,
+            "max_serve_nodes": cluster.max_serve_nodes,
+            "cpu_utilization_percent": cluster.cpu_utilization_percent,
+        },
     }
-    return {"status": "SUCCESS", "results": table_info}
   except Exception as ex:
+    logger.exception("Bigtable metadata tool failed: %s", ex)
     return {
         "status": "ERROR",
-        "error_details": str(ex),
+        "error_details": repr(ex),
     }