auto-retry BQ table-row-count in case it times out

src/metabase/driver/bigquery.clj

@@ -182,12 +182,12 @@
    "TIMESTAMP" parse-timestamp-str})
 
 (def ^:private ^:const query-timeout-error-message "Query timed out.")
-(def ^:private ^:const query-default-timeout-seconds 30)
+(def ^:private ^:const query-default-timeout-seconds 20)
 
 (defn- post-process-native
   ([^QueryResponse response]
    (post-process-native response query-default-timeout-seconds))
-  ([^QueryResponse response, ^Integer timeout-seconds]
+  ([^QueryResponse response, ^Integer timeout-seconds, should-retry?]
    (if-not (.getJobComplete response)
      ;; 99% of the time by the time this is called `.getJobComplete` will return `true`. On the off chance it doesn't, wait a few seconds for the job to finish.
      (do

test/metabase/test/data/bigquery.clj

@@ -90,9 +90,13 @@
   (println (u/format-color 'blue "Created BigQuery table '%s.%s'." dataset-id table-id)))
 
 (defn- table-row-count ^Integer [^String dataset-id, ^String table-id]
-  (first (first (:rows (post-process-native (execute (.query (.jobs bigquery) project-id
-                                                             (doto (QueryRequest.)
-                                                               (.setQuery (format "SELECT COUNT(*) FROM [%s.%s]" dataset-id table-id))))))))))
+  (let [f (fn [] (first (first (:rows (post-process-native (execute (.query (.jobs bigquery) project-id
+                                                                            (doto (QueryRequest.)
+                                                                              (.setQuery (format "SELECT COUNT(*) FROM [%s.%s]" dataset-id table-id))))))))))]
+    ;; automatically retry the query in case it fails / times out
+    (try (f)
+         (catch Throwable _
+           (f)))))
 
 ;; This is a dirty HACK
 (defn- ^DateTime timestamp-korma-form->GoogleDateTime