Fix PostRelease Nightly Snapshot (#33798)

Amar3tto · web-flow · commit fd5e212737af · 2025-01-30T08:48:20.000-05:00
* Update mobile gaming groovy scripts

* Add retry
diff --git a/examples/java/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToBigQuery.java b/examples/java/src/main/java/org/apache/beam/examples/complete/game/utils/WriteToBigQuery.java
@@ -28,6 +28,7 @@
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition;
+import org.apache.beam.sdk.io.gcp.bigquery.InsertRetryPolicy;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -129,7 +130,8 @@ public PDone expand(PCollection<InputT> teamAndScore) {
                 .to(getTable(projectId, datasetId, tableName))
                 .withSchema(getSchema())
                 .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
-                .withWriteDisposition(WriteDisposition.WRITE_APPEND));
+                .withWriteDisposition(WriteDisposition.WRITE_APPEND)
+                .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors()));
     return PDone.in(teamAndScore.getPipeline());
   }
 
diff --git a/release/src/main/groovy/mobilegaming-java-dataflow.groovy b/release/src/main/groovy/mobilegaming-java-dataflow.groovy
@@ -98,15 +98,20 @@ class LeaderBoardRunner {
     def isSuccess = false
     String query_result = ""
     while ((System.currentTimeMillis() - startTime) / 60000 < mobileGamingCommands.EXECUTION_TIMEOUT_IN_MINUTES) {
-      tables = t.run "bq query SELECT table_id FROM ${t.bqDataset()}.__TABLES_SUMMARY__"
-      if (tables.contains("leaderboard_${runner}_user") && tables.contains("leaderboard_${runner}_team")) {
-        query_result = t.run """bq query --batch "SELECT user FROM [${t.gcpProject()}:${
-          t.bqDataset()
-        }.leaderboard_${runner}_user] LIMIT 10\""""
-        if (t.seeAnyOf(mobileGamingCommands.COLORS, query_result)) {
-          isSuccess = true
-          break
+      try {
+        tables = t.run "bq query --use_legacy_sql=false SELECT table_name FROM ${t.bqDataset()}.INFORMATION_SCHEMA.TABLES"
+        if (tables.contains("leaderboard_${runner}_user") && tables.contains("leaderboard_${runner}_team")) {
+          query_result = t.run """bq query --batch "SELECT user FROM [${t.gcpProject()}:${
+            t.bqDataset()
+          }.leaderboard_${runner}_user] LIMIT 10\""""
+          if (t.seeAnyOf(mobileGamingCommands.COLORS, query_result)) {
+            isSuccess = true
+            break
+          }
         }
+      } catch (Exception e) {
+        println "Warning: Exception while checking tables: ${e.message}"
+        println "Retrying..."
       }
       println "Waiting for pipeline to produce more results..."
       sleep(60000) // wait for 1 min
diff --git a/release/src/main/groovy/mobilegaming-java-direct.groovy b/release/src/main/groovy/mobilegaming-java-direct.groovy
@@ -87,13 +87,18 @@ def startTime = System.currentTimeMillis()
 def isSuccess = false
 String query_result = ""
 while((System.currentTimeMillis() - startTime)/60000 < mobileGamingCommands.EXECUTION_TIMEOUT_IN_MINUTES) {
-  tables = t.run "bq query SELECT table_id FROM ${t.bqDataset()}.__TABLES_SUMMARY__"
-  if(tables.contains("leaderboard_${runner}_user") && tables.contains("leaderboard_${runner}_team")){
-    query_result = t.run """bq query --batch "SELECT user FROM [${t.gcpProject()}:${t.bqDataset()}.leaderboard_${runner}_user] LIMIT 10\""""
-    if(t.seeAnyOf(mobileGamingCommands.COLORS, query_result)){
-      isSuccess = true
-      break
+  try {
+    tables = t.run "bq query --use_legacy_sql=false SELECT table_name FROM ${t.bqDataset()}.INFORMATION_SCHEMA.TABLES"
+    if(tables.contains("leaderboard_${runner}_user") && tables.contains("leaderboard_${runner}_team")) {
+      query_result = t.run """bq query --batch "SELECT user FROM [${t.gcpProject()}.${t.bqDataset()}.leaderboard_${runner}_user] LIMIT 10\""""
+      if(t.seeAnyOf(mobileGamingCommands.COLORS, query_result)){
+        isSuccess = true
+        break
+      }
     }
+  } catch (Exception e) {
+    println "Warning: Exception while checking tables: ${e.message}"
+    println "Retrying..."
   }
   println "Waiting for pipeline to produce more results..."
   sleep(60000) // wait for 1 min
