Fix flaking database-schema-filtering-test for bigquery (#36900)

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

@@ -63,23 +63,29 @@
 ;;; |                                                      Sync                                                      |
 ;;; +----------------------------------------------------------------------------------------------------------------+
 
+(defn- list-datasets
+  "Fetch all datasets given database `details`, applying dataset filters if specified."
+  [{:keys [project-id dataset-filters-type dataset-filters-patterns] :as details}]
+  (let [client (database-details->client details)
+        project-id (or project-id (bigquery.common/database-details->credential-project-id details))
+        datasets (.listDatasets client project-id (u/varargs BigQuery$DatasetListOption))
+        inclusion-patterns (when (= "inclusion" dataset-filters-type) dataset-filters-patterns)
+        exclusion-patterns (when (= "exclusion" dataset-filters-type) dataset-filters-patterns)]
+    (for [^Dataset dataset (.iterateAll datasets)
+          :let [^DatasetId dataset-id (.. dataset getDatasetId)]
+          :when (driver.s/include-schema? inclusion-patterns
+                                          exclusion-patterns
+                                          (.getDataset dataset-id))]
+      dataset-id)))
+
 (defn- list-tables
   "Fetch all tables (new pages are loaded automatically by the API)."
-  (^Iterable [database-details]
-   (list-tables database-details {:validate-dataset? false}))
-  (^Iterable [{:keys [project-id dataset-filters-type dataset-filters-patterns] :as details} {:keys [validate-dataset?]}]
+  (^Iterable [details]
+   (list-tables details {:validate-dataset? false}))
+  (^Iterable [details {:keys [validate-dataset?]}]
    (let [client (database-details->client details)
-         project-id (or project-id (bigquery.common/database-details->credential-project-id details))
-         datasets (.listDatasets client project-id (u/varargs BigQuery$DatasetListOption))
-         inclusion-patterns (when (= "inclusion" dataset-filters-type) dataset-filters-patterns)
-         exclusion-patterns (when (= "exclusion" dataset-filters-type) dataset-filters-patterns)
-         dataset-iter (for [^Dataset dataset (.iterateAll datasets)
-                            :let [^DatasetId dataset-id (.. dataset getDatasetId)]
-                            :when (driver.s/include-schema? inclusion-patterns
-                                                            exclusion-patterns
-                                                            (.getDataset dataset-id))]
-                        dataset-id)]
-     (when (and (not= dataset-filters-type "all")
+         dataset-iter (list-datasets details)]
+     (when (and (not= (:dataset-filters-type details) "all")
                 validate-dataset?
                 (zero? (count dataset-iter)))
        (throw (ex-info (tru "Looks like we cannot find any matching datasets.")

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

@@ -500,30 +500,51 @@
 (defn- sync-and-assert-filtered-tables [database assert-table-fn]
   (t2.with-temp/with-temp [Database db-filtered database]
     (sync/sync-database! db-filtered {:scan :schema})
-    (doseq [table (t2/select-one Table :db_id (u/the-id db-filtered))]
-      (assert-table-fn table))))
+    (let [tables (t2/select Table :db_id (u/the-id db-filtered))]
+      (assert (not-empty tables))
+      (doseq [table tables]
+        (assert-table-fn table)))))
 
 (deftest dataset-filtering-test
   (mt/test-driver :bigquery-cloud-sdk
     (testing "Filtering BigQuery connections for datasets works as expected"
-      (testing " with an inclusion filter"
-        (sync-and-assert-filtered-tables {:name    "BigQuery Test DB with dataset inclusion filters"
-                                          :engine  :bigquery-cloud-sdk
-                                          :details (-> (mt/db)
-                                                       :details
-                                                       (assoc :dataset-filters-type "inclusion"
-                                                              :dataset-filters-patterns "a*,t*"))}
-                                         (fn [{dataset-id :schema}]
-                                           (is (not (contains? #{\a \t} (first dataset-id)))))))
-      (testing " with an exclusion filter"
-        (sync-and-assert-filtered-tables {:name    "BigQuery Test DB with dataset exclusion filters"
-                                          :engine  :bigquery-cloud-sdk
-                                          :details (-> (mt/db)
-                                                       :details
-                                                       (assoc :dataset-filters-type "exclusion"
-                                                              :dataset-filters-patterns "v*"))}
-          (fn [{dataset-id :schema}]
-            (is (not= \v (first dataset-id)))))))))
+      (mt/db) ;; force the creation of one test dataset
+      (mt/dataset avian-singles
+        (mt/db) ;; force the creation of another test dataset
+        (let [;; This test is implemented in this way to avoid having to create new datasets, and to avoid
+              ;; syncing most of the tables in the test DB.
+              datasets (#'bigquery/list-datasets (-> (mt/db)
+                                                     :details
+                                                     (dissoc :dataset-filters-type
+                                                             :dataset-filters-patterns)))
+              dataset-ids (map #(.getDataset %) datasets)
+              ;; get the first 4 characters of each dataset-id
+              prefixes (->> dataset-ids
+                            (map (fn [dataset-id]
+                                   (apply str (take 4 dataset-id)))))
+              ;; inclusion-patterns selects the first dataset
+              ;; exclusion-patterns excludes almost everything else
+              include-prefix (first prefixes)
+              inclusion-patterns (str include-prefix "*")
+              exclusion-patterns (str/join "," (map #(str % "*") (set (rest prefixes))))]
+          (testing " with an inclusion filter"
+            (sync-and-assert-filtered-tables {:name    "BigQuery Test DB with dataset inclusion filters"
+                                              :engine  :bigquery-cloud-sdk
+                                              :details (-> (mt/db)
+                                                           :details
+                                                           (assoc :dataset-filters-type "inclusion"
+                                                                  :dataset-filters-patterns inclusion-patterns))}
+                                             (fn [{dataset-id :schema}]
+                                               (is (str/starts-with? dataset-id include-prefix)))))
+          (testing " with an exclusion filter"
+            (sync-and-assert-filtered-tables {:name    "BigQuery Test DB with dataset exclusion filters"
+                                              :engine  :bigquery-cloud-sdk
+                                              :details (-> (mt/db)
+                                                           :details
+                                                           (assoc :dataset-filters-type "exclusion"
+                                                                  :dataset-filters-patterns exclusion-patterns))}
+                                             (fn [{dataset-id :schema}]
+                                               (is (str/starts-with? dataset-id include-prefix))))))))))
 
 (deftest normalize-away-dataset-id-test
   (mt/test-driver :bigquery-cloud-sdk

test/metabase/driver/sql_jdbc/sync/describe_database_test.clj

@@ -154,7 +154,9 @@
          driver)))
 
 (deftest database-schema-filtering-test
-  (mt/test-drivers (schema-filtering-drivers)
+  ;; BigQuery is tested separately in `metabase.driver.bigquery-cloud-sdk-test/dataset-filtering-test`, because
+  ;; otherwise this test takes too long and flakes intermittently
+  (mt/test-drivers (disj (schema-filtering-drivers) :bigquery-cloud-sdk)
     (let [driver             (driver.u/database->driver (mt/db))
           schema-filter-prop (find-schema-filters-prop driver)
           filter-type-prop   (keyword (str (:name schema-filter-prop) "-type"))