Bigquery - Ignore retry on query cancellation (#36625)

* Dont retry cancelled queries on Bigquery

* Add test coverage for ignoring retry on cancellation

* Cancel running BQ jobs on cancellation request

* Fix indentation

* Throw exception with cancelled flag

* Join lines

* Change unique query string

* Fix parentheses

Co-authored-by: metamben <103100869+metamben@users.noreply.github.com>

* Change e to ex-data

* Set namespace explictly in mock exception

* Rename query identification string

* Use bigquery namespace

Co-authored-by: metamben <103100869+metamben@users.noreply.github.com>

* Remove query cancellation code

* Remove comment

---------

Co-authored-by: metamben <103100869+metamben@users.noreply.github.com>

modules/drivers/bigquery-cloud-sdk/src/metabase/driver/bigquery_cloud_sdk.clj

@@ -301,8 +301,7 @@
                 @res-fut)))))
       @res-fut)
     (catch java.util.concurrent.CancellationException _e
-      ;; trying to deref the value after the future has been cancelled
-      (throw (ex-info (tru "Query cancelled") {:sql sql :parameters parameters})))
+      (throw (ex-info (tru "Query cancelled") {:sql sql :parameters parameters ::cancelled? true})))
     (catch BigQueryException e
       (if (.isRetryable e)
         (throw (ex-info (tru "BigQueryException executing query")
@@ -373,7 +372,8 @@
       (thunk)
       (catch Throwable e
         (let [ex-data (u/all-ex-data e)]
-          (if (or (:retryable? e) (not (qp.error-type/client-error? (:type ex-data))))
+          (if (and (not (::cancelled? ex-data))
+                   (or (:retryable? ex-data) (not (qp.error-type/client-error? (:type ex-data)))))
             (thunk)
             (throw e)))))))
 

modules/drivers/bigquery-cloud-sdk/test/metabase/driver/bigquery_cloud_sdk_test.clj

@@ -439,6 +439,29 @@
           ;; make sure that the fake exception was thrown, and thus the query execution was retried
           (is (true? @fake-execute-called)))))))
 
+(deftest not-retry-cancellation-exception-test
+  (mt/test-driver :bigquery-cloud-sdk
+    (let [fake-execute-called (atom false)
+          orig-fn        @#'bigquery/execute-bigquery]
+      (testing "Should not retry query on cancellation"
+        (with-redefs [bigquery/execute-bigquery (fn [^BigQuery client ^String sql parameters _ _]
+                                                  ;; We only want to simulate exception on the query that we're testing and not on possible db setup queries
+                                                  (if (and (re-find #"notRetryCancellationExceptionTest" sql) (not @fake-execute-called))
+                                                    (do (reset! fake-execute-called true)
+                                                        ;; Simulate a cancellation happening
+                                                        (throw (ex-info "Query cancelled" {::bigquery/cancelled? true})))
+                                                    (orig-fn client sql parameters nil nil)))]
+          (try
+            (qp/process-query {:native {:query "SELECT CURRENT_TIMESTAMP() AS notRetryCancellationExceptionTest"} :database (mt/id)
+                               :type     :native})
+            ;; If no exception is thrown, then the test should fail
+            (is false "Query should have failed")
+            (catch clojure.lang.ExceptionInfo e
+              ;; Verify exception as expected
+              (is (= "Query cancelled" (.getMessage e)))
+              ;; make sure that the fake exception was thrown
+              (is (true? @fake-execute-called)))))))))
+
 (deftest query-cancel-test
   (mt/test-driver :bigquery-cloud-sdk
     (testing "BigQuery queries can be canceled successfully"