Test data loading improvements [ci drivers]

modules/drivers/druid/test/metabase/test/data/druid.clj

@@ -30,13 +30,14 @@
 ;;; Generating Data File
 
 (defn- flattened-test-data []
-  (let [dbdef    (tx/flatten-dbdef defs/test-data "checkins")
+  (let [dbdef    (tx/get-dataset-definition
+                  (tx/flattened-dataset-definition defs/test-data "checkins"))
         tabledef (first (:table-definitions dbdef))]
     (->> (:rows tabledef)
          (map (partial zipmap (map :field-name (:field-definitions tabledef))))
          (map-indexed (fn [i row]
                         (assoc row :id (inc i))))
-         (sort-by (u/rpartial get "date")))))
+         (sort-by #(get % "date")))))
 
 (defn- write-dbdef-to-json [db-def filename]
   (io/delete-file filename :silently)

modules/drivers/mongo/test/metabase/driver/mongo_test.clj

@@ -126,12 +126,12 @@
   (driver/describe-table :mongo (data/db) (Table (data/id :venues))))
 
 ;; Make sure that all-NULL columns work and are synced correctly (#6875)
-(tx/def-database-definition ^:private all-null-columns
+(tx/defdataset ^:private all-null-columns
   [["bird_species"
-     [{:field-name "name", :base-type :type/Text}
-      {:field-name "favorite_snack", :base-type :type/Text}]
-     [["House Finch" nil]
-      ["Mourning Dove" nil]]]])
+    [{:field-name "name", :base-type :type/Text}
+     {:field-name "favorite_snack", :base-type :type/Text}]
+    [["House Finch" nil]
+     ["Mourning Dove" nil]]]])
 
 (datasets/expect-with-driver :mongo
   [{:name "_id",            :database_type "java.lang.Long",   :base_type :type/Integer, :special_type :type/PK}
@@ -198,7 +198,7 @@
 
 
 ;;; Check that we support Mongo BSON ID and can filter by it (#1367)
-(tx/def-database-definition ^:private with-bson-ids
+(tx/defdataset ^:private with-bson-ids
   [["birds"
      [{:field-name "name", :base-type :type/Text}
       {:field-name "bird_id", :base-type :type/MongoBSONID}]

modules/drivers/sqlserver/test/metabase/driver/sqlserver_test.clj

@@ -16,16 +16,16 @@
              [util :as tu :refer [obj->json->obj]]]
             [metabase.test.data
              [datasets :as datasets]
-             [interface :as tx :refer [def-database-definition]]]))
+             [interface :as tx]]))
 
 ;;; -------------------------------------------------- VARCHAR(MAX) --------------------------------------------------
 
 ;; Make sure something long doesn't come back as some weird type like `ClobImpl`
-(def ^:private ^:const a-gene
+(def ^:private a-gene
   "Really long string representing a gene like \"GGAGCACCTCCACAAGTGCAGGCTATCCTGTCGAGTAAGGCCT...\""
   (apply str (repeatedly 1000 (partial rand-nth [\A \G \C \T]))))
 
-(def-database-definition ^:private ^:const genetic-data
+(tx/defdataset ^:private genetic-data
   [["genetic-data"
      [{:field-name "gene", :base-type {:native "VARCHAR(MAX)"}}]
      [[a-gene]]]])

src/metabase/db.clj

@@ -466,7 +466,7 @@
       (u/with-us-locale
         (verify-db-connection db-details)
         (run-schema-migrations! auto-migrate db-details)
-        (create-connection-pool! (du/profile (jdbc-details db-details)))
+        (create-connection-pool! (jdbc-details db-details))
         (run-data-migrations!)
         (reset! db-setup-finished? true))))
   nil)

src/metabase/query_processor.clj

@@ -233,9 +233,11 @@
 (def ^:private default-pipeline (qp-pipeline execute-query))
 
 (def ^:private QueryResponse
-  (s/cond-pre
-   ManyToManyChannel
-   {:status (s/enum :completed :failed :canceled), s/Any s/Any}))
+  (s/named
+   (s/cond-pre
+    ManyToManyChannel
+    {:status (s/enum :completed :failed :canceled), s/Any s/Any})
+   "Valid query response (core.async channel or map with :status)"))
 
 (s/defn process-query :- QueryResponse
   "Process an MBQL query. This is the main entrypoint to the magical realm of the Query Processor. Returns a

test/metabase/api/dataset_test.clj

@@ -203,7 +203,7 @@
    ["3" "2014-09-15" "" "8" "56"]
    ["4" "2014-03-11" "" "5" "4"]
    ["5" "2013-05-05" "" "3" "49"]]
-  (data/with-db (data/get-or-create-database! defs/test-data-with-null-date-checkins)
+  (data/dataset defs/test-data-with-null-date-checkins
     (let [result ((test-users/user->client :rasta) :post 200 "dataset/csv" :query
                   (json/generate-string (data/mbql-query checkins)))]
       (take 5 (parse-and-sort-csv result)))))

test/metabase/api/table_test.clj

@@ -23,7 +23,6 @@
              [data :as data]
              [util :as tu :refer [match-$]]]
             [metabase.test.data
-             [dataset-definitions :as defs]
              [datasets :as datasets]
              [users :refer [user->client]]]
             [metabase.test.mock.util :as mutil]
@@ -205,24 +204,6 @@
             :fields_hash  $}))
   ((user->client :rasta) :get 200 (format "table/%d/query_metadata" (data/id :categories))))
 
-
-(def ^:private user-last-login-date-strs
-  "In an effort to be really annoying, the date strings returned by the API are different on Circle than they are
-   locally. Generate strings like '2014-01-01' at runtime so we get matching values."
-  (let [format-inst (fn [^java.util.Date inst]
-                      (format "%d-%02d-%02d"
-                              (+ (.getYear inst) 1900)
-                              (+ (.getMonth inst) 1)
-                              (.getDate inst)))]
-    (->> (defs/field-values defs/test-data-map "users" "last_login")
-         (map format-inst)
-         set
-         sort
-         vec)))
-
-(def ^:private user-full-names
-  (defs/field-values defs/test-data-map "users" "name"))
-
 ;;; GET api/table/:id/query_metadata?include_sensitive_fields
 ;; Make sure that getting the User table *does* include info about the password field, but not actual values
 ;; themselves

test/metabase/driver/mysql_test.clj

@@ -17,7 +17,7 @@
              [util :as tu]]
             [metabase.test.data
              [datasets :refer [expect-with-driver]]
-             [interface :as tx :refer [def-database-definition]]]
+             [interface :as tx]]
             [metabase.test.util.timezone :as tu.tz]
             [metabase.util.date :as du]
             [toucan.db :as db]
@@ -54,7 +54,7 @@
 ;; Test how TINYINT(1) columns are interpreted. By default, they should be interpreted as integers, but with the
 ;; correct additional options, we should be able to change that -- see
 ;; https://github.com/metabase/metabase/issues/3506
-(def-database-definition ^:private ^:const tiny-int-ones
+(tx/defdataset ^:private tiny-int-ones
   [["number-of-cans"
      [{:field-name "thing",          :base-type :type/Text}
       {:field-name "number-of-cans", :base-type {:native "tinyint(1)"}}]
@@ -71,7 +71,7 @@
   #{{:name "number-of-cans", :base_type :type/Boolean, :special_type :type/Category}
     {:name "id",             :base_type :type/Integer, :special_type :type/PK}
     {:name "thing",          :base_type :type/Text,    :special_type :type/Category}}
-  (data/with-temp-db [db tiny-int-ones]
+  (data/with-db-for-dataset [db tiny-int-ones]
     (db->fields db)))
 
 ;; if someone says specifies `tinyInt1isBit=false`, it should come back as a number instead
@@ -79,7 +79,7 @@
   #{{:name "number-of-cans", :base_type :type/Integer, :special_type :type/Quantity}
     {:name "id",             :base_type :type/Integer, :special_type :type/PK}
     {:name "thing",          :base_type :type/Text,    :special_type :type/Category}}
-  (data/with-temp-db [db tiny-int-ones]
+  (data/with-db-for-dataset [db tiny-int-ones]
     (tt/with-temp Database [db {:engine "mysql"
                                 :details (assoc (:details db)
                                            :additional-options "tinyInt1isBit=false")}]

test/metabase/driver/postgres_test.clj

@@ -64,9 +64,9 @@
 ;; Verify that we identify JSON columns and mark metadata properly during sync
 (datasets/expect-with-driver :postgres
   :type/SerializedJSON
-  (data/with-temp-db
+  (data/with-db-for-dataset
     [_
-     (tx/create-database-definition "Postgres with a JSON Field"
+     (tx/dataset-definition "Postgres with a JSON Field"
        ["venues"
         [{:field-name "address", :base-type {:native "json"}}]
         [[(hsql/raw "to_json('{\"street\": \"431 Natoma\", \"city\": \"San Francisco\", \"state\": \"CA\", \"zip\": 94103}'::text)")]]])]
@@ -74,7 +74,7 @@
 
 
 ;;; # UUID Support
-(tx/def-database-definition ^:private with-uuid
+(tx/defdataset ^:private with-uuid
   [["users"
     [{:field-name "user_id", :base-type :type/UUID}]
     [[#uuid "4f01dcfd-13f7-430c-8e6f-e505c0851027"]
@@ -125,7 +125,7 @@
 
 
 ;; Make sure that Tables / Fields with dots in their names get escaped properly
-(tx/def-database-definition ^:private dots-in-names
+(tx/defdataset ^:private dots-in-names
   [["objects.stuff"
     [{:field-name "dotted.name", :base-type :type/Text}]
     [["toucan_cage"]
@@ -143,7 +143,7 @@
 
 
 ;; Make sure that duplicate column names (e.g. caused by using a FK) still return both columns
-(tx/def-database-definition ^:private duplicate-names
+(tx/defdataset ^:private duplicate-names
   [["birds"
     [{:field-name "name", :base-type :type/Text}]
     [["Rasta"]
@@ -163,7 +163,7 @@
 
 
 ;;; Check support for `inet` columns
-(tx/def-database-definition ^:private ip-addresses
+(tx/defdataset ^:private ip-addresses
   [["addresses"
     [{:field-name "ip", :base-type {:native "inet"}}]
     [[(hsql/raw "'192.168.1.1'::inet")]

test/metabase/query_processor/async_test.clj

@@ -9,20 +9,21 @@
 
 ;; running a query async should give you the same results as running that query synchronously
 (let [query
-      {:database (data/id)
-       :type     :query
-       :query    {:source-table (data/id :venues)
-                  :fields       [[:field-id (data/id :venues :name)]]
-                  :limit        5}}
+      (delay
+       {:database (data/id)
+        :type     :query
+        :query    {:source-table (data/id :venues)
+                   :fields       [[:field-id (data/id :venues :name)]]
+                   :limit        5}})
       ;; Metadata checksum might be encrypted if a encryption key is set on this system (to make it hard for bad
       ;; actors to forge one) in which case the checksums won't be equal.
       maybe-decrypt-checksum
       #(some-> % (update-in [:data :results_metadata :checksum] encrypt/maybe-decrypt))]
   (expect
     (maybe-decrypt-checksum
-     (qp/process-query query))
+     (qp/process-query @query))
     (maybe-decrypt-checksum
-     (tu.async/with-open-channels [result-chan (qp.async/process-query query)]
+     (tu.async/with-open-channels [result-chan (qp.async/process-query @query)]
        (first (a/alts!! [result-chan (a/timeout 1000)]))))))
 
 (expect

test/metabase/query_processor_test.clj

@@ -335,7 +335,8 @@
   [results]
   (when (= (:status results) :failed)
     (println "Error running query:" (u/pprint-to-str 'red results))
-    (throw (ex-info (:error results) results)))
+    (throw (ex-info (str (or (:error results) "Error running query"))
+             (if (map? results) results {:results results}))))
   (:data results))
 
 (defn rows

test/metabase/query_processor_test/breakout_test.clj

@@ -14,9 +14,7 @@
             [metabase.test
              [data :as data]
              [util :as tu]]
-            [metabase.test.data
-             [dataset-definitions :as defs]
-             [datasets :as datasets]]
+            [metabase.test.data.datasets :as datasets]
             [metabase.test.util.log :as tu.log]
             [toucan.db :as db]
             [toucan.util.test :as tt]))
@@ -107,7 +105,7 @@
    :native_form true}
   (data/with-data
     (fn []
-      (let [venue-names (defs/field-values defs/test-data-map "categories" "name")]
+      (let [venue-names (data/dataset-field-values "categories" "name")]
         [(db/insert! Dimension {:field_id (data/id :venues :category_id)
                                 :name     "Foo"
                                 :type     :internal})

test/metabase/query_processor_test/remapping_test.clj

@@ -2,7 +2,7 @@
   "Tests for the remapping results"
   (:require [metabase
              [query-processor :as qp]
-             [query-processor-test :refer :all]]
+             [query-processor-test :as qp.test :refer :all]]
             [metabase.models
              [dimension :refer [Dimension]]
              [field :refer [Field]]]
@@ -16,10 +16,10 @@
             [toucan.db :as db]))
 
 (qp-expect-with-all-drivers
-  {:rows        [["20th Century Cafe" 12 "Café Sweets"]
-                 ["25°" 11 "Café"]
-                 ["33 Taps" 7 "Beer Garden"]
-                 ["800 Degrees Neapolitan Pizzeria" 58 "Ramen"]]
+  {:rows        [["20th Century Cafe"               12 "Café"]
+                 ["25°"                             11 "Burger"]
+                 ["33 Taps"                          7 "Bar"]
+                 ["800 Degrees Neapolitan Pizzeria" 58 "Pizza"]]
    :columns     [(data/format-name "name")
                  (data/format-name "category_id")
                  "Foo"]
@@ -33,8 +33,8 @@
            {:fields   [$name $category_id]
             :order-by [[:asc $name]]
             :limit    4})
-         booleanize-native-form
-         (format-rows-by [str int str])
+         qp.test/booleanize-native-form
+         (qp.test/format-rows-by [str int str])
          tu/round-fingerprint-cols)))
 
 (defn- select-columns
@@ -78,8 +78,8 @@
     (->> (data/run-mbql-query venues
            {:order-by [[:asc $name]]
             :limit    4})
-         booleanize-native-form
-         (format-rows-by [int str int double double int str])
+         qp.test/booleanize-native-form
+         (qp.test/format-rows-by [int str int double double int str])
          (select-columns (set (map data/format-name ["name" "price" "name_2"])))
          tu/round-fingerprint-cols
          data)))
@@ -107,8 +107,8 @@
            {:fields   [$name $price $category_id]
             :order-by [[:asc $name]]
             :limit    4})
-         booleanize-native-form
-         (format-rows-by [str int str str])
+         qp.test/booleanize-native-form
+         (qp.test/format-rows-by [str int str str])
          (select-columns (set (map data/format-name ["name" "price" "name_2"])))
          tu/round-fingerprint-cols
          :data)))

test/metabase/query_processor_test/time_field_test.clj

@@ -87,10 +87,7 @@
     [[1 "Plato Yeshua" "08:30:00.000Z"]
      [4 "Simcha Yan" "08:30:00.000Z"]])
   (tu/with-temporary-setting-values [report-timezone "America/Los_Angeles"]
-    (time-query {:filter (vec (cons
-                               :between
-                               (cons
-                                $last_login_time
-                                (if (qpt/supports-report-timezone? driver/*driver*)
-                                  ["08:00:00" "09:00:00"]
-                                  ["08:00:00-00:00" "09:00:00-00:00"]))))})))
+    (time-query {:filter (into [:between $last_login_time]
+                               (if (qpt/supports-report-timezone? driver/*driver*)
+                                 ["08:00:00" "09:00:00"]
+                                 ["08:00:00-00:00" "09:00:00-00:00"]))})))

test/metabase/query_processor_test/timezones_test.clj

@@ -12,16 +12,10 @@
             [metabase.test.data
              [dataset-definitions :as defs]
              [datasets :refer [expect-with-drivers]]
-             [interface :as i]
              [sql :as sql.tx]]
             [toucan.db :as db]))
 
 (defn- call-with-timezones-db [f]
-  ;; Does the database exist?
-  (when-not (i/metabase-instance defs/test-data-with-timezones driver/*driver*)
-    ;; The database doesn't exist, so we need to create it
-    (data/get-or-create-database! defs/test-data-with-timezones))
-  ;; The database can now be used in tests
   (data/with-db (data/get-or-create-database! defs/test-data-with-timezones)
     (f)))
 

test/metabase/sync/sync_metadata/comments_test.clj

@@ -21,7 +21,7 @@
     (set (map (partial into {}) (db/select ['Field :name :description] :table_id [:in table-ids])))))
 
 ;; test basic field comments sync
-(tx/def-database-definition ^:const ^:private basic-field-comments
+(tx/defdataset ^:private basic-field-comments
   [["basic_field_comments"
     [{:field-name "with_comment", :base-type :type/Text, :field-comment "comment"}
      {:field-name "no_comment", :base-type :type/Text}]
@@ -31,11 +31,11 @@
   #{{:name (data/format-name "id"), :description nil}
     {:name (data/format-name "with_comment"), :description "comment"}
     {:name (data/format-name "no_comment"), :description nil}}
-  (data/with-temp-db [db basic-field-comments]
+  (data/with-db-for-dataset [db basic-field-comments]
     (db->fields db)))
 
 ;; test changing the description in metabase db so we can check it is not overwritten by comment in source db when resyncing
-(tx/def-database-definition ^:const ^:private update-desc
+(tx/defdataset ^:private update-desc
   [["update_desc"
     [{:field-name "updated_desc", :base-type :type/Text, :field-comment "original comment"}]
     [["foo"]]]])
@@ -43,7 +43,7 @@
 (datasets/expect-with-drivers #{:h2 :postgres}
   #{{:name (data/format-name "id"), :description nil}
     {:name (data/format-name "updated_desc"), :description "updated description"}}
-  (data/with-temp-db [db update-desc]
+  (data/with-db-for-dataset [db update-desc]
     ;; change the description in metabase while the source table comment remains the same
     (db/update-where! Field {:id (data/id "update_desc" "updated_desc")}, :description "updated description")
     ;; now sync the DB again, this should NOT overwrite the manually updated description
@@ -51,7 +51,7 @@
     (db->fields db)))
 
 ;; test adding a comment to the source data that was initially empty, so we can check that the resync picks it up
-(tx/def-database-definition ^:const ^:private comment-after-sync
+(tx/defdataset ^:private comment-after-sync
   [["comment_after_sync"
     [{:field-name "comment_after_sync", :base-type :type/Text}]
     [["foo"]]]])
@@ -59,14 +59,20 @@
 (datasets/expect-with-drivers #{:h2 :postgres}
   #{{:name (data/format-name "id"), :description nil}
     {:name (data/format-name "comment_after_sync"), :description "added comment"}}
-  (data/with-temp-db [db comment-after-sync]
+  (data/with-db-for-dataset [db comment-after-sync]
     ;; modify the source DB to add the comment and resync. The easiest way to do this is just destroy the entire DB
     ;; and re-create a modified version. As such, let the SQL JDBC driver know the DB is being "modified" so it can
     ;; destroy its current connection pool
     (driver/notify-database-updated driver/*driver* db)
-    (let [modified-dbdef (assoc-in comment-after-sync
-                                   [:table-definitions 0 :field-definitions 0 :field-comment]
-                                   "added comment")]
+    (let [modified-dbdef (update
+                          comment-after-sync
+                          :table-definitions
+                          (fn [[tabledef]]
+                            [(update
+                              tabledef
+                              :field-definitions
+                              (fn [[fielddef]]
+                                [(assoc fielddef :field-comment "added comment")]))]))]
       (tx/create-db! driver/*driver* modified-dbdef))
     (sync/sync-table! (Table (data/id "comment_after_sync")))
     (db->fields db)))
@@ -87,13 +93,13 @@
 ;; test basic comments on table
 (datasets/expect-with-drivers #{:h2 :postgres}
   #{{:name (data/format-name "table_with_comment"), :description "table comment"}}
-  (data/with-temp-db [db (basic-table "table_with_comment" "table comment")]
+  (data/with-db-for-dataset [db (basic-table "table_with_comment" "table comment")]
     (db->tables db)))
 
 ;; test changing the description in metabase on table to check it is not overwritten by comment in source db when resyncing
 (datasets/expect-with-drivers #{:h2 :postgres}
   #{{:name (data/format-name "table_with_updated_desc"), :description "updated table description"}}
-  (data/with-temp-db [db (basic-table "table_with_updated_desc" "table comment")]
+  (data/with-db-for-dataset [db (basic-table "table_with_updated_desc" "table comment")]
     ;; change the description in metabase while the source table comment remains the same
     (db/update-where! Table {:id (data/id "table_with_updated_desc")}, :description "updated table description")
     ;; now sync the DB again, this should NOT overwrite the manually updated description
@@ -103,7 +109,7 @@
 ;; test adding a comment to the source table that was initially empty, so we can check that the resync picks it up
 (datasets/expect-with-drivers #{:h2 :postgres}
   #{{:name (data/format-name "table_with_comment_after_sync"), :description "added comment"}}
-  (data/with-temp-db [db (basic-table "table_with_comment_after_sync" nil)]
+  (data/with-db-for-dataset [db (basic-table "table_with_comment_after_sync" nil)]
     ;; modify the source DB to add the comment and resync
     (driver/notify-database-updated driver/*driver* db)
     (tx/create-db! driver/*driver* (basic-table "table_with_comment_after_sync" "added comment"))

test/metabase/sync/sync_metadata/tables_test.clj

@@ -11,7 +11,7 @@
             [toucan.db :as db]
             [toucan.util.test :as tt]))
 
-(tx/def-database-definition ^:private db-with-some-cruft
+(tx/defdataset ^:private db-with-some-cruft
   [["acquired_toucans"
      [{:field-name "species",              :base-type :type/Text}
       {:field-name "cam_has_acquired_one", :base-type :type/Boolean}]

test/metabase/test/data.clj

 (ns metabase.test.data
-  "Code related to creating and deleting test databases + datasets."
+  "Super-useful test utility functions.
+
+  Basic way stuff in here, which you'll see everywhere in the tests, is:
+
+  1. Get the DB you're currently testing by calling `db`. Get IDs of the DB or of its Fields and Tables in that DB by
+     calling `id`.
+
+      (data/db)                 ; -> Get current test DB
+      (data/id)                 ; -> Get ID of current test DB
+      (data/id :table)          ; -> Get ID of Table named `table` in current test DB
+      (data/id :table :field)   ; -> Get ID of Field named `field` belonging to Table `table` in current test DB
+
+     Normally this database is the `test-data` database for the current driver, and is created the first time `db` or
+     `id` is called.
+
+  2. Bind the current driver with `driver/with-driver`. Defaults to `:h2`
+
+       (driver/with-driver :postgres
+         (data/id))
+       ;; -> Get ID of Postgres `test-data` database, creating it if needed
+
+  3. Bind a different database for use with for `db` and `id` functions with `with-db`.
+
+      (data/with-db [db some-database]
+        (data/id :table :field)))
+       ;; -> Return ID of Field named `field` in Table `table` in `some-db`
+
+  4. You can use helper macros like `$ids` to replace symbols starting with `$` (for Fields) or `$$` (for Tables) with
+     calls to `id` in a form:
+
+      ($ids {:source-table $$venues, :fields [$venues.name]})
+      ;; -> {:source-table (data/id :venues), :fields [(data/id :venues :name)]}
+
+     (There are several variations of this macro; see documentation below for more details.)"
   (:require [cheshire.core :as json]
             [clojure
              [string :as str]
@@ -22,18 +55,19 @@
              [dataset-definitions :as defs]
              [interface :as tx]]
             [metabase.test.util.timezone :as tu.tz]
-            [toucan.db :as db])
-  (:import [metabase.test.data.interface DatabaseDefinition TableDefinition]))
+            [schema.core :as s]
+            [toucan.db :as db]))
 
 (declare get-or-create-database!)
 
+
 ;;; ------------------------------------------ Dataset-Independent Data Fns ------------------------------------------
 
 ;; These functions offer a generic way to get bits of info like Table + Field IDs from any of our many driver/dataset
 ;; combos.
 
 (defn get-or-create-test-data-db!
-  "Get or create the Test Data database for DRIVER, which defaults to driver/`*driver*`."
+  "Get or create the Test Data database for `driver`, which defaults to `driver/*driver*`, or `:h2` if that is unbound."
   ([]       (get-or-create-test-data-db! (tx/driver)))
   ([driver] (get-or-create-database! driver defs/test-data)))
 
@@ -54,8 +88,7 @@
     (f)))
 
 (defmacro with-db
-  "Run body with DB as the current database.
-   Calls to `db` and `id` use this value."
+  "Run body with `db` as the current database. Calls to `db` and `id` use this value."
   [db & body]
   `(do-with-db ~db (fn [] ~@body)))
 
@@ -98,23 +131,23 @@
   to `id`.
 
   Optionally wraps IDs in `:field-id` or `:fk->` clauses as appropriate; this defaults to true."
-  [table-name body & {:keys [wrap-field-ids?], :or {wrap-field-ids? true}}]
+  [table-name-or-nil body & {:keys [wrap-field-ids?], :or {wrap-field-ids? true}}]
   (walk/postwalk
    (fn [form]
      (cond
        (not (symbol? form))
        form
 
-       (= form '$$table)
-       `(id ~(keyword table-name))
+       (and table-name-or-nil (= form '$$table))
+       `(id ~(keyword table-name-or-nil))
 
        (str/starts-with? form "$$")
-       (let [table-name (str/replace form #"^\$\$" "")]
-         `(id ~(keyword table-name)))
+       (let [table-name-or-nil (str/replace form #"^\$\$" "")]
+         `(id ~(keyword table-name-or-nil)))
 
        (str/starts-with? form "$")
        (let [field-name (str/replace form #"^\$" "")]
-         (token->id-call wrap-field-ids? table-name field-name))
+         (token->id-call wrap-field-ids? table-name-or-nil field-name))
 
        :else
        form))
@@ -125,7 +158,7 @@
   With no `.` delimiters, it is assumed we're referring to a Field belonging to `table-name`, which is passed implicitly
   as the first arg. With one or more `.` delimiters, no implicit `table-name` arg is passed to `id`:
 
-    $venue_id      -> (id :sightings :venue_id) ; TABLE-NAME is implicit first arg
+    $venue_id      -> (id :sightings :venue_id) ; `table-name` is implicit first arg
     $cities.id     -> (id :cities :id)          ; specify non-default Table
 
   Use `$$table` to refer to the table itself.
@@ -141,13 +174,16 @@
     ($ids [venues {:wrap-field-ids? true}]
       $category_id->categories.name)
     ;; -> [:fk-> [:field-id (id :venues :category_id(] [:field-id (id :categories :name)]]"
-  {:arglists '([table & body] [[table {:keys [wrap-field-ids?]}] & body]), :style/indent 1}
-  [table-and-options & body]
-  (let [[table-name options] (if (sequential? table-and-options)
-                               table-and-options
-                               [table-and-options])]
-    (m/mapply $->id (keyword table-name) `(do ~@body) (merge {:wrap-field-ids? false}
-                                                             options))))
+  {:arglists '([form] [table & body] [[table {:keys [wrap-field-ids?], :or {wrap-field-ids? false}}] & body])}
+  (^{:style/indent 0} [form]
+   `($ids nil ~form))
+
+  (^{:style/indent 1} [table-and-options & body]
+   (let [[table-name options] (if (sequential? table-and-options)
+                                table-and-options
+                                [table-and-options])]
+     (m/mapply $->id (keyword table-name) `(do ~@body) (merge {:wrap-field-ids? false}
+                                                              options)))))
 
 (declare id)
 
@@ -266,14 +302,14 @@
 
 (defn- add-extra-metadata!
   "Add extra metadata like Field base-type, etc."
-  [database-definition db]
-  (doseq [^TableDefinition table-definition (:table-definitions database-definition)]
-    (let [table-name (:table-name table-definition)
-          table      (delay (or  (tx/metabase-instance table-definition db)
-                                 (throw (Exception. (format "Table '%s' not loaded from definiton:\n%s\nFound:\n%s"
-                                                            table-name
-                                                            (u/pprint-to-str (dissoc table-definition :rows))
-                                                            (u/pprint-to-str (db/select [Table :schema :name], :db_id (:id db))))))))]
+  [{:keys [table-definitions], :as database-definition} db]
+  {:pre [(seq table-definitions)]}
+  (doseq [{:keys [table-name], :as table-definition} table-definitions]
+    (let [table (delay (or (tx/metabase-instance table-definition db)
+                           (throw (Exception. (format "Table '%s' not loaded from definiton:\n%s\nFound:\n%s"
+                                                      table-name
+                                                      (u/pprint-to-str (dissoc table-definition :rows))
+                                                      (u/pprint-to-str (db/select [Table :schema :name], :db_id (:id db))))))))]
       (doseq [{:keys [field-name visibility-type special-type], :as field-definition} (:field-definitions table-definition)]
         (let [field (delay (or (tx/metabase-instance field-definition @table)
                                (throw (Exception. (format "Field '%s' not loaded from definition:\n"
@@ -286,7 +322,8 @@
             (log/debug (format "SET SPECIAL TYPE %s.%s -> %s" table-name field-name special-type))
             (db/update! Field (:id @field) :special_type (u/keyword->qualified-name special-type))))))))
 
-(defn- create-database! [{:keys [database-name], :as database-definition} driver]
+(defn- create-database! [driver {:keys [database-name], :as database-definition}]
+  {:pre [(seq database-name)]}
   ;; Create the database and load its data
   ;; ALWAYS CREATE DATABASE AND LOAD DATA AS UTC! Unless you like broken tests
   (tu.tz/with-jvm-tz "UTC"
@@ -299,10 +336,14 @@
     ;; sync newly added DB
     (sync/sync-database! db)
     ;; add extra metadata for fields
-    (add-extra-metadata! database-definition db)
+    (try
+      (add-extra-metadata! database-definition db)
+      (catch Throwable e
+        (println "Error adding extra metadata:" e)))
     ;; make sure we're returing an up-to-date copy of the DB
     (Database (u/get-id db))))
 
+
 (defonce ^:private ^{:arglists '([driver]), :doc "We'll have a very bad time if any sort of test runs that calls
   `data/db` for the first time calls it multiple times in parallel -- for example my Oracle test that runs 30 sync
   calls at the same time to make sure nothing explodes and cursors aren't leaked. To make sure this doesn't happen
@@ -318,28 +359,38 @@
            (swap! locks update driver #(or % (Object.)))
            (@locks driver)))))))
 
-(defn get-or-create-database!
-  "Create DBMS database associated with DATABASE-DEFINITION, create corresponding Metabase
-  `Databases`/`Tables`/`Fields`, and sync the `Database`. DRIVER should be an object that implements
-  `IDriverTestExtensions`; it defaults to the value returned by the method `driver` for the current
-  dataset (driver/`*driver*`), which is H2 by default."
-  ([database-definition]
-   (get-or-create-database! (tx/driver) database-definition))
-  ([driver database-definition]
-   (let [driver (tx/the-driver-with-test-extensions driver)]
+(defmulti get-or-create-database!
+  "Create DBMS database associated with `database-definition`, create corresponding Metabase Databases/Tables/Fields,
+  and sync the Database. `driver` is a keyword name of a driver that implements test extension methods (as defined in
+  the `metabase.test.data.interface` namespace); `driver` defaults to `driver/*driver*` if bound, or `:h2` if not.
+  `database-definition` is anything that implements the `tx/get-database-definition` method."
+  {:arglists '([database-definition] [driver database-definition])}
+  (fn
+    ([_]
+     (tx/dispatch-on-driver-with-test-extensions (tx/driver)))
+    ([driver _]
+     (tx/dispatch-on-driver-with-test-extensions driver)))
+  :hierarchy #'driver/hierarchy)
+
+(defmethod get-or-create-database! :default
+  ([dbdef]
+   (get-or-create-database! (tx/driver) dbdef))
+
+  ([driver dbdef]
+   (let [dbdef (tx/get-dataset-definition dbdef)]
      (or
-      (tx/metabase-instance database-definition driver)
+      (tx/metabase-instance dbdef driver)
       (locking (driver->create-database-lock driver)
         (or
-         (tx/metabase-instance database-definition driver)
-         (create-database! database-definition driver)))))))
+         (tx/metabase-instance dbdef driver)
+         (create-database! driver dbdef)))))))
 
 
-(defn do-with-temp-db
-  "Execute F with DBDEF loaded as the current dataset. F takes a single argument, the `DatabaseInstance` that was
-  loaded and synced from DBDEF."
-  [^DatabaseDefinition dbdef, f]
-  (let [dbdef  (tx/map->DatabaseDefinition dbdef)]
+(s/defn do-with-db-for-dataset
+  "Execute `f` with `dbdef` loaded as the current dataset. `f` takes a single argument, the DatabaseInstance that was
+  loaded and synced from `dbdef`."
+  [dataset-definition, f :- (s/pred fn?)]
+  (let [dbdef (tx/get-dataset-definition dataset-definition)]
     (binding [db/*disable-db-logging* true]
       (let [db (get-or-create-database! (tx/driver) dbdef)]
         (assert db)
@@ -348,11 +399,11 @@
           (f db))))))
 
 
-(defmacro with-temp-db
-  "Load and sync DATABASE-DEFINITION with DRIVER and execute BODY with the newly created `Database` bound to
-  DB-BINDING, and make it the current database for `metabase.test.data` functions like `id`.
+(defmacro with-db-for-dataset
+  "Load and sync `database-definition` with `driver` and execute `body` with the newly created Database bound to
+  `db-binding`, and make it the current database for `metabase.test.data` functions like `id`.
 
-     (with-temp-db [db tupac-sightings]
+     (with-db-for-dataset [db tupac-sightings]
        (driver/process-quiery {:database (:id db)
                                :type     :query
                                :query    {:source-table (:id &events)
@@ -361,29 +412,29 @@
 
   A given Database is only created once per run of the test suite, and is automatically destroyed at the conclusion
   of the suite."
-  [[db-binding, ^DatabaseDefinition database-definition] & body]
-  `(do-with-temp-db ~database-definition
+  [[db-binding dataset-def] & body]
+  `(do-with-db-for-dataset ~dataset-def
      (fn [~db-binding]
        ~@body)))
 
-(defn resolve-dbdef [symb]
-  @(or (resolve symb)
+(defn resolve-dbdef [namespace-symb symb]
+  @(or (ns-resolve namespace-symb symb)
        (ns-resolve 'metabase.test.data.dataset-definitions symb)
-       (throw (Exception. (format "Dataset definition not found: '%s' or 'metabase.test.data.dataset-definitions/%s'"
-                                  symb symb)))))
+       (throw (Exception. (format "Dataset definition not found: '%s/%s' or 'metabase.test.data.dataset-definitions/%s'"
+                                  namespace-symb symb symb)))))
 
 (defmacro dataset
-  "Load and sync a temporary `Database` defined by DATASET, make it the current DB (for `metabase.test.data` functions
-  like `id`), and execute BODY.
+  "Load and sync a temporary Database defined by `dataset`, make it the current DB (for `metabase.test.data` functions
+  like `id` and `db`), and execute `body`.
 
-  Like `with-temp-db`, but takes an unquoted symbol naming a `DatabaseDefinition` rather than the dbef itself.
-  DATASET is optionally namespace-qualified; if not, `metabase.test.data.dataset-definitions` is assumed.
+  Like `with-db-for-dataset`, but takes an unquoted symbol naming a DatabaseDefinition rather than the dbef itself. `dataset`
+  is optionally namespace-qualified; if not, `metabase.test.data.dataset-definitions` is assumed.
 
      (dataset sad-toucan-incidents
        ...)"
   {:style/indent 1}
   [dataset & body]
-  `(with-temp-db [~'_ (resolve-dbdef '~dataset)]
+  `(with-db-for-dataset [~'_ (resolve-dbdef '~(ns-name *ns*) '~dataset)]
      ~@body))
 
 (defn- delete-model-instance!
@@ -392,7 +443,7 @@
   (db/delete! (-> instance name symbol) :id id))
 
 (defn call-with-data
-  "Takes a thunk `DATA-LOAD-FN` that returns a seq of toucan model instances that will be deleted after `BODY-FN`
+  "Takes a thunk `data-load-fn` that returns a seq of toucan model instances that will be deleted after `body-fn`
   finishes"
   [data-load-fn body-fn]
   (let [result-instances (data-load-fn)]
@@ -402,30 +453,51 @@
         (doseq [instance result-instances]
           (delete-model-instance! instance))))))
 
-(defmacro with-data [data-load-fn & body]
+(defmacro with-data
+  "Calls `data-load-fn` to create a sequence of objects, then runs `body`; finally, deletes the objects."
+  [data-load-fn & body]
   `(call-with-data ~data-load-fn (fn [] ~@body)))
 
-(def ^:private venue-categories
-  (map vector (defs/field-values defs/test-data-map "categories" "name")))
+(defn dataset-field-values
+  "Get all the values for a field in a `dataset-definition`.
 
+    (dataset-field-values \"categories\" \"name\") ; -> [\"African\" \"American\" \"Artisan\" ...]"
+  ([table-name field-name]
+   (dataset-field-values defs/test-data table-name field-name))
+
+  ([dataset-definition table-name field-name]
+   (some
+    (fn [{:keys [field-definitions rows], :as tabledef}]
+      (when (= table-name (:table-name tabledef))
+        (some
+         (fn [[i fielddef]]
+           (when (= field-name (:field-name fielddef))
+             (map #(nth % i) rows)))
+         (m/indexed field-definitions))))
+    (:table-definitions (tx/get-dataset-definition dataset-definition)))))
+
+(def ^:private category-names
+  (delay (vec (dataset-field-values "categories" "name"))))
+
+;; TODO - you should always call these functions with the `with-data` macro. We should enforce this
 (defn create-venue-category-remapping
-  "Returns a thunk that adds an internal remapping for category_id in the venues table aliased as `REMAPPING-NAME`.
+  "Returns a thunk that adds an internal remapping for category_id in the venues table aliased as `remapping-name`.
   Can be used in a `with-data` invocation."
   [remapping-name]
   (fn []
     [(db/insert! Dimension {:field_id (id :venues :category_id)
-                            :name remapping-name
-                            :type :internal})
-     (db/insert! FieldValues {:field_id (id :venues :category_id)
-                              :values (json/generate-string (range 0 (count venue-categories)))
-                              :human_readable_values (json/generate-string (map first venue-categories))})]))
+                            :name     remapping-name
+                            :type     :internal})
+     (db/insert! FieldValues {:field_id              (id :venues :category_id)
+                              :values                (json/generate-string (range 1 (inc (count @category-names))))
+                              :human_readable_values (json/generate-string @category-names)})]))
 
 (defn create-venue-category-fk-remapping
-  "Returns a thunk that adds a FK remapping for category_id in the venues table aliased as `REMAPPING-NAME`. Can be
+  "Returns a thunk that adds a FK remapping for category_id in the venues table aliased as `remapping-name`. Can be
   used in a `with-data` invocation."
   [remapping-name]
   (fn []
-    [(db/insert! Dimension {:field_id (id :venues :category_id)
-                            :name remapping-name
-                            :type :external
+    [(db/insert! Dimension {:field_id                (id :venues :category_id)
+                            :name                    remapping-name
+                            :type                    :external
                             :human_readable_field_id (id :categories :name)})]))

test/metabase/test/data/dataset_definitions.clj

 (ns metabase.test.data.dataset-definitions
-  "Definitions of various datasets for use in tests with `with-temp-db`."
-  (:require [metabase.test.data.interface :as di])
+  "Definitions of various datasets for use in tests with `data/dataset` or `data/with-db-for-dataset`."
+  (:require [medley.core :as m]
+            [metabase.test.data.interface :as tx])
   (:import java.sql.Time
            [java.util Calendar TimeZone]))
 
-;; ## Datasets
+;;; +----------------------------------------------------------------------------------------------------------------+
+;;; |                                                Various Datasets                                                |
+;;; +----------------------------------------------------------------------------------------------------------------+
 
-;; The O.G. "Test Database" dataset
-(di/def-database-definition-edn test-data)
+(tx/defdataset-edn test-data
+  "The O.G. \"Test Database\" dataset.")
 
-;; Times when the Toucan cried
-(di/def-database-definition-edn sad-toucan-incidents)
+ (tx/defdataset-edn sad-toucan-incidents
+   "Times when the Toucan cried")
 
-;; Places, times, and circumstances where Tupac was sighted. Sighting timestamps are UNIX Timestamps in seconds
-(di/def-database-definition-edn tupac-sightings)
+(tx/defdataset-edn tupac-sightings
+  "Places, times, and circumstances where Tupac was sighted. Sighting timestamps are UNIX Timestamps in seconds")
 
-;; Dataset with nested columns, for testing a MongoDB-style database
-(di/def-database-definition-edn geographical-tips)
+(tx/defdataset-edn geographical-tips
+  "Dataset with nested columns, for testing a MongoDB-style database")
 
-;; A very tiny dataset with a list of places and a booleans
-(di/def-database-definition-edn places-cam-likes)
+(tx/defdataset-edn places-cam-likes
+  "A very tiny dataset with a list of places and a booleans")
 
-;; A small dataset with users and a set of messages between them. Each message has *2* foreign keys to user --
-;; sender and receiver -- allowing us to test situations where multiple joins for a *single* table should occur.
-(di/def-database-definition-edn avian-singles)
+(tx/defdataset-edn avian-singles
+  "A small dataset with users and a set of messages between them. Each message has *2* foreign keys to user -- sender
+  and receiver -- allowing us to test situations where multiple joins for a *single* table should occur.")
 
-;; A small dataset that includes an integer column with some NULL and ZERO values, meant for testing things like
-;; expressions to make sure they behave correctly
-;;
-;; As an added "bonus" this dataset has a table with a name in a slash in it, so the driver will need to support that
-;; correctly in order for this to work!
-(di/def-database-definition-edn daily-bird-counts)
+(tx/defdataset-edn daily-bird-counts
+  "A small dataset that includes an integer column with some NULL and ZERO values, meant for testing things like
+  expressions to make sure they behave correctly.
+
+  As an added bonus this dataset has a table with a name in a slash in it, so the driver will need to support that
+  correctly in order for this to work!")
+
+(tx/defdataset-edn office-checkins
+  "A small dataset that includes TIMESTAMP dates. People who stopped by the Metabase office and the time they did so.")
+
+(tx/defdataset-edn bird-flocks
+  "A small dataset with birds and the flocks they belong to (a many-to-one relationship). Some birds belong to no
+  flocks, and one flock has no birds, so this is useful for testing behavior of various types of joins. (`flock_id` is
+  not explicitly marked as a foreign key, because the test dataset syntax does not yet have a way to support nullable
+  foreign keys.)")
 
-;; A small dataset that includes TIMESTAMP dates. People who stopped by the Metabase office and the time they did so.
-(di/def-database-definition-edn office-checkins)
 
 (defn- calendar-with-fields ^Calendar [date & fields]
-  (let [cal-from-date  (doto (Calendar/getInstance (TimeZone/getTimeZone "UTC"))
-                         (.setTime date))
-        blank-calendar (doto (.clone cal-from-date)
-                         .clear)]
+  (let [^Calendar cal-from-date  (doto (Calendar/getInstance (TimeZone/getTimeZone "UTC"))
+                                   (.setTime date))
+        ^Calendar blank-calendar (doto ^Calendar (.clone cal-from-date)
+                                   .clear)]
     (doseq [field fields]
       (.set blank-calendar field (.get cal-from-date field)))
     blank-calendar))
@@ -57,66 +67,72 @@
   [date]
   (Time. (.getTimeInMillis (calendar-with-fields date Calendar/HOUR_OF_DAY Calendar/MINUTE Calendar/SECOND))))
 
-(di/def-database-definition test-data-with-time
-  (di/update-table-def "users"
-                       (fn [table-def]
-                         [(first table-def)
-                          {:field-name "last_login_date", :base-type :type/Date}
-                          {:field-name "last_login_time", :base-type :type/Time}
-                          (peek table-def)])
-                       (fn [rows]
-                         (mapv (fn [[username last-login password-text]]
-                                 [username (date-only last-login) (time-only last-login) password-text])
-                               rows))
-                       (for [[table-name :as orig-def] (di/slurp-edn-table-def "test-data")
-                             :when (= table-name "users")]
-                         orig-def)))
-
-(di/def-database-definition test-data-with-null-date-checkins
-  (di/update-table-def "checkins"
-                       #(vec (concat % [{:field-name "null_only_date" :base-type :type/Date}]))
-                       (fn [rows]
-                         (mapv #(conj % nil) rows))
-                       (di/slurp-edn-table-def "test-data")))
-
-(di/def-database-definition test-data-with-timezones
-  (di/update-table-def "users"
-                       (fn [table-def]
-                         [(first table-def)
-                          {:field-name "last_login", :base-type :type/DateTimeWithTZ}
-                          (peek table-def)])
-                       identity
-                       (di/slurp-edn-table-def "test-data")))
-
-(def test-data-map
-  "Converts data from `test-data` to a map of maps like the following:
-
-   {<table-name> [{<field-name> <field value> ...}]."
-  (reduce (fn [acc {:keys [table-name field-definitions rows]}]
-            (let [field-names (mapv :field-name field-definitions)]
-              (assoc acc table-name
-                     (for [row rows]
-                       (zipmap field-names row)))))
-          {} (:table-definitions test-data)))
-
-(defn field-values
-  "Returns the field values for the given `TABLE` and `COLUMN` found
-  in the data-map `M`."
-  [m table column]
-  (mapv #(get % column) (get m table)))
-
-;; Takes the `test-data` dataset and adds a `created_by` column to the users table that is self referencing
-(di/def-database-definition test-data-self-referencing-user
-  (di/update-table-def "users"
-                       (fn [table-def]
-                         (conj table-def {:field-name "created_by", :base-type :type/Integer, :fk :users}))
-                       (fn [rows]
-                         (mapv (fn [[username last-login password-text] idx]
-                                 [username last-login password-text (if (= 1 idx)
-                                                                      idx
-                                                                      (dec idx))])
-                               rows
-                               (iterate inc 1)))
-                       (for [[table-name :as orig-def] (di/slurp-edn-table-def "test-data")
-                             :when (= table-name "users")]
-                         orig-def)))
+
+(defonce ^{:doc "The main `test-data` dataset, but only the `users` table, and with `last_login_date` and
+  `last_login_time` instead of `last_login`."}
+  test-data-with-time
+  (tx/transformed-dataset-definition "test-data-with-time" test-data
+    (tx/transform-dataset-only-tables "users")
+    (tx/transform-dataset-update-table "users"
+      :table
+      (fn [tabledef]
+        (update
+         tabledef
+         :field-definitions
+         (fn [[name-field-def _ password-field-def]]
+           [name-field-def
+            (tx/map->FieldDefinition {:field-name "last_login_date", :base-type :type/Date})
+            (tx/map->FieldDefinition {:field-name "last_login_time", :base-type :type/Time})
+            password-field-def])))
+      :rows
+      (fn [rows]
+        (for [[username last-login password-text] rows]
+          [username (date-only last-login) (time-only last-login) password-text])))))
+
+(defonce ^{:doc "The main `test-data` dataset, with an additional (all-null) `null_only_date` Field."}
+  test-data-with-null-date-checkins
+  (tx/transformed-dataset-definition "test-data-with-null-date-checkins" test-data
+    (tx/transform-dataset-update-table "checkins"
+      :table
+      (fn [tabledef]
+        (update
+         tabledef
+         :field-definitions
+         concat
+         [(tx/map->FieldDefinition {:field-name "null_only_date" :base-type :type/Date})]))
+      :rows
+      (fn [rows]
+        (for [row rows]
+          (concat row [nil]))))))
+
+(defonce ^{:doc "The main `test-data` dataset, but `last_login` has a base type of `:type/DateTimeWithTZ`."}
+  test-data-with-timezones
+  (tx/transformed-dataset-definition "test-data-with-timezones" test-data
+    (tx/transform-dataset-update-table "users"
+      :table
+      (fn [tabledef]
+        (update
+         tabledef
+         :field-definitions
+         (fn [[name-field-def _ password-field-def]]
+           [name-field-def
+            (tx/map->FieldDefinition {:field-name "last_login", :base-type :type/DateTimeWithTZ})
+            password-field-def]))))))
+
+(defonce ^{:doc "The usual `test-data` dataset, but only the `users` table; adds a `created_by` column to the users
+  table that is self referencing."}
+  test-data-self-referencing-user
+  (tx/transformed-dataset-definition "test-data-self-referencing-user" test-data
+    (tx/transform-dataset-only-tables "users")
+    (tx/transform-dataset-update-table "users"
+      :table
+      (fn [tabledef]
+        (update tabledef :field-definitions concat [(tx/map->FieldDefinition
+                                                     {:field-name "created_by", :base-type :type/Integer, :fk :users})]))
+      ;; created_by = user.id - 1, except for User 1, who was created by himself (?)
+      :rows
+      (fn [rows]
+        (for [[idx [username last-login password-text]] (m/indexed rows)]
+          [username last-login password-text (if (zero? idx)
+                                               1
+                                               idx)])))))