Fix xlml_to_buganizer.py cannot query two clusters with the same names in two locations (GoogleCloudPlatform#954)

* support multi location for the same cluster in the same project

* update

* remove useless code

* update review

Author: severus-ho

dags/dashboard/xlml_to_buganizer.py

@@ -79,6 +79,102 @@ def build_malfunction_row(
   ]
 
 
+def insert_cluster_status_lists(
+    credentials, status_list, project_name, cluster_name, location, now_utc
+):
+  """Insert both cluster + nodepool status into list"""
+  try:
+    #  Cannot query cluster info without location
+    if not location:
+      raise NotFound(
+          "The requested resource can not be found on the server without location."
+      )
+    info = get_cluster_status(credentials, project_name, location, cluster_name)
+    cluster_status = info["status"]
+    cluster_status_message = info["status_message"]
+    mal_node_pools = [
+        node_pool
+        for node_pool in info["node_pools"]
+        if node_pool["status"]
+        not in [
+            Status.RUNNING.value,
+            Status.PROVISIONING.value,
+            Status.STOPPING.value,
+        ]
+    ]
+    is_cluster_malfunction = (
+        cluster_status != Status.RUNNING.value
+        and cluster_status != Status.RECONCILING.value
+    )
+    malfunction_node_pools_exist = len(mal_node_pools) > 0
+
+    if is_cluster_malfunction and malfunction_node_pools_exist:
+      status_list.append(
+          build_malfunction_row(
+              issue_type=IssueType.CLUSTER_AND_NODE_POOL,
+              proj=project_name,
+              cluster_name=cluster_name,
+              cluster_location=location,
+              cluster_status=cluster_status,
+              cluster_status_message=cluster_status_message,
+              node_pools=mal_node_pools,
+              now_utc=now_utc,
+          )
+      )
+    elif is_cluster_malfunction:
+      status_list.append(
+          build_malfunction_row(
+              issue_type=IssueType.CLUSTER,
+              proj=project_name,
+              cluster_location=location,
+              cluster_name=cluster_name,
+              cluster_status=cluster_status,
+              cluster_status_message=cluster_status_message,
+              now_utc=now_utc,
+          )
+      )
+    elif malfunction_node_pools_exist:
+      status_list.append(
+          build_malfunction_row(
+              issue_type=IssueType.NODE_POOL,
+              proj=project_name,
+              cluster_name=cluster_name,
+              cluster_location=location,
+              cluster_status=cluster_status,
+              cluster_status_message=cluster_status_message,
+              node_pools=mal_node_pools,
+              now_utc=now_utc,
+          )
+      )
+  except NotFound:
+    status_list.append(
+        build_malfunction_row(
+            issue_type=IssueType.CLUSTER,
+            proj=project_name,
+            cluster_name=cluster_name,
+            cluster_location=UNKNOW_LOCATION,
+            cluster_status=Status.NOT_EXIST.value,
+            cluster_status_message=NOT_FOUND_MSG,
+            now_utc=now_utc,
+        )
+    )
+  except Exception as e:
+    logging.info(
+        f"Error fetching details for {project_name}/{cluster_name}: {e}"
+    )
+    status_list.append(
+        build_malfunction_row(
+            issue_type=IssueType.CLUSTER,
+            proj=project_name,
+            cluster_name=cluster_name,
+            cluster_location=UNKNOW_LOCATION,
+            cluster_status=Status.ERROR.value,
+            cluster_status_message=str(e),
+            now_utc=now_utc,
+        )
+    )
+
+
 def print_failed_cluster_info(cluster_status_rows):
   """Log failed cluster info for debugging"""
   result_rows_list = []
@@ -115,7 +211,7 @@ def insert_gspread_rows(rows: List[List[str]], sheet_id: str):
         insert_data_option="INSERT_ROWS",
         value_input_option="RAW",
     )
-    logging.info(f"Successfully appended rows to Google Sheet.")
+    logging.info(f"Successfully appended {len(rows)} rows to Google Sheet.")
   except Exception as e:
     logging.error(f"An error occurred while writing to Google Sheet: {e}")
 
@@ -130,7 +226,7 @@ def get_clusters_from_view(credential) -> List[Any]:
   dataset_id = context["params"]["source_bq_dataset_id"] or DEFAULT_DATASET_ID
   client = bigquery.Client(project=project_id, credentials=credential)
   query = f"""
-        SELECT project_name, cluster_name
+        SELECT project_name, cluster_name, region
         FROM `{project_id}.{dataset_id}.cluster_view`
     """
   rows = list(client.query(query).result())
@@ -141,14 +237,6 @@ def get_clusters_from_view(credential) -> List[Any]:
 # ================================================================
 # Step 4: GKE Functions
 # ================================================================
-def list_clusters_by_project(credentials, project_id) -> List[Any]:
-  """List all clusters in a given GCP project"""
-  client = container_v1.ClusterManagerClient(credentials=credentials)
-  parent = f"projects/{project_id}/locations/-"
-  response = client.list_clusters(request={"parent": parent})
-  return response.clusters if response and response.clusters else []
-
-
 def get_cluster_status(
     credentials, project_id, location, cluster_name
 ) -> Dict[str, Any]:
@@ -180,7 +268,7 @@ def get_cluster_status(
 
   return {
       "project_id": project_id,
-      "region": location,
+      "location": location,
       "cluster_name": cluster_name,
       "status": container_v1.Cluster.Status(cluster.status).name
       if cluster.status
@@ -193,144 +281,25 @@ def get_cluster_status(
 # ================================================================
 # Step 5: Workflow Functions
 # ================================================================
-def fetch_clusters_list(credential):
+def fetch_clusters(credential) -> List[str]:
   clusters = get_clusters_from_view(credential=credential)
   if not clusters:
     logging.info("No rows found in view.")
     return []
   return clusters
 
 
-def fetch_clusters_location(credentials, clusters) -> Dict[str, Any]:
-  """Map project::cluster_name -> location"""
-  cluster_locations = {}
-  projects = set(cluster.project_name for cluster in clusters)
-  logging.info(f"Fetching clusters from {len(projects)} distinct projects...")
-  for idx, project in enumerate(projects, 1):
-    try:
-      clusters = list_clusters_by_project(
-          credentials=credentials, project_id=project
-      )
-      logging.info(
-          f"[{idx}/{len(projects)}] {project}: {len(clusters)} clusters"
-      )
-      for cluster in clusters:
-        key = f"{project}::{cluster.name}"
-        cluster_locations[key] = cluster.location
-    except Exception as e:
-      logging.info(f"Error listing clusters for {project}: {e}")
-  return cluster_locations
-
-
-def insert_cluster_status_lists(
-    credentials, status_list, project_name, location, cluster_name, now_utc
-):
-  """Insert both cluster + nodepool status into list"""
-  try:
-    #  Cannot query cluster info without location
-    if not location:
-      raise NotFound("The requested resource was not found on the server.")
-    info = get_cluster_status(credentials, project_name, location, cluster_name)
-    cluster_status = info["status"]
-    cluster_status_message = info["status_message"]
-    mal_node_pools = [
-        node_pool
-        for node_pool in info["node_pools"]
-        if node_pool["status"] != Status.RUNNING.value
-        and node_pool["status"] != Status.PROVISIONING.value
-        and node_pool["status"] != Status.STOPPING.value
-    ]
-    is_cluster_malfunction = (
-        cluster_status != Status.RUNNING.value
-        and cluster_status != Status.RECONCILING.value
-    )
-    malfunction_node_pools_exist = len(mal_node_pools) > 0
-
-    if is_cluster_malfunction and malfunction_node_pools_exist:
-      status_list.append(
-          build_malfunction_row(
-              issue_type=IssueType.CLUSTER_AND_NODE_POOL,
-              proj=project_name,
-              cluster_name=cluster_name,
-              cluster_location=location,
-              cluster_status=cluster_status,
-              cluster_status_message=cluster_status_message,
-              node_pools=mal_node_pools,
-              now_utc=now_utc,
-          )
-      )
-    elif is_cluster_malfunction:
-      status_list.append(
-          build_malfunction_row(
-              issue_type=IssueType.CLUSTER,
-              proj=project_name,
-              cluster_location=location,
-              cluster_name=cluster_name,
-              cluster_status=cluster_status,
-              cluster_status_message=cluster_status_message,
-              now_utc=now_utc,
-          )
-      )
-    elif malfunction_node_pools_exist:
-      status_list.append(
-          build_malfunction_row(
-              issue_type=IssueType.NODE_POOL,
-              proj=project_name,
-              cluster_name=cluster_name,
-              cluster_location=location,
-              cluster_status=cluster_status,
-              cluster_status_message=cluster_status_message,
-              node_pools=mal_node_pools,
-              now_utc=now_utc,
-          )
-      )
-  except NotFound:
-    status_list.append(
-        build_malfunction_row(
-            issue_type=IssueType.CLUSTER,
-            proj=project_name,
-            cluster_name=cluster_name,
-            cluster_location=UNKNOW_LOCATION,
-            cluster_status=Status.NOT_EXIST.value,
-            cluster_status_message=NOT_FOUND_MSG,
-            now_utc=now_utc,
-        )
-    )
-  except Exception as e:
-    logging.info(
-        f"Error fetching details for {project_name}/{cluster_name}: {e}"
-    )
-    status_list.append(
-        build_malfunction_row(
-            issue_type=IssueType.CLUSTER,
-            proj=project_name,
-            cluster_name=cluster_name,
-            cluster_location=UNKNOW_LOCATION,
-            cluster_status=Status.ERROR.value,
-            cluster_status_message=str(e),
-            now_utc=now_utc,
-        )
-    )
-
-
-def fetch_clusters_status(
-    credentials, clusters_list: List[Any], cluster_locations: Dict[str, Any]
-) -> List[Any]:
+def fetch_clusters_status(credentials, clusters: List[Any]) -> List[Any]:
   """Fetch status for all clusters & node pools"""
   cluster_status_rows = []
   now_utc = datetime.now(timezone.utc).isoformat()
-  for cluster in clusters_list:
-    project_name = cluster.project_name
-    cluster_name = cluster.cluster_name
-    key = f"{project_name}::{cluster_name}"
-    location = cluster_locations.get(key)
-    #  cluster_status_rows is parameter for output
+  for cluster in clusters:
     insert_cluster_status_lists(
         credentials,
         cluster_status_rows,
-        project_name,
-        location,
-        cluster_name,
+        cluster.project_name,
+        cluster.cluster_name,
+        cluster.region,
         now_utc,
     )
   print_failed_cluster_info(cluster_status_rows)
@@ -345,13 +314,8 @@ def pull_clusters_status() -> List[Any]:
   credentials, _ = google.auth.default(
       scopes=["https://www.googleapis.com/auth/cloud-platform"]
   )
-  target_clusters_list = fetch_clusters_list(credentials)
-  target_clusters_locations = fetch_clusters_location(
-      credentials, target_clusters_list
-  )
-  return fetch_clusters_status(
-      credentials, target_clusters_list, target_clusters_locations
-  )
+  target_clusters = fetch_clusters(credentials)
+  return fetch_clusters_status(credentials, target_clusters)
 
 
 @task