TableDefinition can define a custom PK (#39038)

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

@@ -128,19 +128,16 @@
 ;; Customize the create table table to include the S3 location
 ;; TODO: Specify a unique location each time
 (defmethod sql.tx/create-table-sql :athena
-  [driver {:keys [database-name]} {:keys [table-name], :as tabledef}]
-  (let [field-definitions (cons
-                           {:field-name "id", :base-type :type/Integer, :semantic-type :type/PK}
-                           (:field-definitions tabledef))
-        fields            (->> field-definitions
-                               (map (fn [{:keys [field-name base-type]}]
-                                      (format "`%s` %s"
-                                              (ddl.i/format-name driver field-name)
-                                              (if (map? base-type)
-                                                (:native base-type)
-                                                (sql.tx/field-base-type->sql-type driver base-type)))))
-                               (interpose ", ")
-                               str/join)]
+  [driver {:keys [database-name]} {:keys [table-name field-definitions], :as _tabledef}]
+  (let [fields (->> field-definitions
+                    (map (fn [{:keys [field-name base-type]}]
+                           (format "`%s` %s"
+                                   (ddl.i/format-name driver field-name)
+                                   (if (map? base-type)
+                                     (:native base-type)
+                                     (sql.tx/field-base-type->sql-type driver base-type)))))
+                    (interpose ", ")
+                    str/join)]
     ;; ICEBERG tables do what we want, and dropping them causes the data to disappear; dropping a normal non-ICEBERG
     ;; table doesn't delete data, so if you recreate it you'll have duplicate rows. 'normal' tables do not support
     ;; `DELETE .. FROM`, either, so there's no way to fix them here.
@@ -204,7 +201,7 @@
             ;; This tells Athena to convert `timestamp with time zone` literals to `timestamp` because otherwise it gets
             ;; very fussy! See [[athena/*loading-data*]] for more info.
             athena/*loading-data*  true]
-    (apply load-data/load-data-add-ids-chunked! args)))
+    (apply load-data/load-data-maybe-add-ids-chunked! args)))
 
 (defn- server-connection-details []
   (tx/dbdef->connection-details :athena :server nil))

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

@@ -248,12 +248,12 @@
          (is (true? (t2/select-one-fn :database_indexed :model/Field (mt/id :singly-index :indexed))))
          (is (false? (t2/select-one-fn :database_indexed :model/Field (mt/id :singly-index :not-indexed)))))
 
-        (testing "compount index"
-          (mongo.connection/with-mongo-database [db (mt/db)]
-            (mongo.util/create-index (mongo.util/collection db "compound-index") (array-map "first" 1 "second" 1)))
-          (sync/sync-database! (mt/db))
-          (is (true? (t2/select-one-fn :database_indexed :model/Field (mt/id :compound-index :first))))
-          (is (false? (t2/select-one-fn :database_indexed :model/Field (mt/id :compound-index :second)))))
+       (testing "compount index"
+         (mongo.connection/with-mongo-database [db (mt/db)]
+           (mongo.util/create-index (mongo.util/collection db "compound-index") (array-map "first" 1 "second" 1)))
+         (sync/sync-database! (mt/db))
+         (is (true? (t2/select-one-fn :database_indexed :model/Field (mt/id :compound-index :first))))
+         (is (false? (t2/select-one-fn :database_indexed :model/Field (mt/id :compound-index :second)))))
 
        (testing "multi key index"
          (mongo.connection/with-mongo-database [db (mt/db)]
@@ -671,14 +671,14 @@
   (mt/test-driver :mongo
     (testing "Negative values in versionArray are ignored (#29678)"
       (with-redefs [mongo.util/run-command (constantly {"version" "4.0.28-23"
-                                                       "versionArray" [4 0 29 -100]})]
+                                                        "versionArray" [4 0 29 -100]})]
         (is (= {:version "4.0.28-23"
                 :semantic-version [4 0 29]}
                (driver/dbms-version :mongo (mt/db))))))
 
     (testing "Any values after rubbish in versionArray are ignored"
       (with-redefs [mongo.util/run-command (constantly {"version" "4.0.28-23"
-                                                       "versionArray" [4 0 "NaN" 29]})]
+                                                        "versionArray" [4 0 "NaN" 29]})]
         (is (= {:version "4.0.28-23"
                 :semantic-version [4 0]}
                (driver/dbms-version :mongo (mt/db))))))))

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

@@ -118,7 +118,7 @@
 
 (defmethod load-data/load-data! :oracle
   [driver dbdef tabledef]
-  (load-data/load-data-add-ids-chunked! driver dbdef tabledef))
+  (load-data/load-data-maybe-add-ids-chunked! driver dbdef tabledef))
 
 ;; Oracle has weird syntax for inserting multiple rows, it looks like
 ;;

modules/drivers/presto-jdbc/test/metabase/test/data/presto_jdbc.clj

@@ -162,7 +162,8 @@
 
 (deftest ^:parallel create-table-sql-test
   (testing "Make sure logic to strip out NOT NULL and PRIMARY KEY stuff works as expected"
-    (let [db-def    (tx/get-dataset-definition defs/test-data)
+    (let [db-def    (update (tx/get-dataset-definition defs/test-data)
+                            :table-definitions (partial #'ddl/add-pks-if-needed :presto-jdbc))
           table-def (-> db-def :table-definitions second)]
       (is (= "CREATE TABLE \"test_data\".\"default\".\"test_data_categories\" (\"id\" INTEGER, \"name\" VARCHAR) ;"
              (sql.tx/create-table-sql :presto-jdbc db-def table-def))))))

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

@@ -192,7 +192,7 @@
 
 (defmethod load-data/load-data! :snowflake
   [& args]
-  (apply load-data/load-data-add-ids-chunked! args))
+  (apply load-data/load-data-maybe-add-ids-chunked! args))
 
 (defmethod tx/aggregate-column-info :snowflake
   ([driver ag-type]

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

@@ -100,15 +100,13 @@
            (throw e)))))))
 
 (defmethod load-data/load-data! :sparksql [& args]
-  (apply load-data/load-data-add-ids! args))
+  (apply load-data/load-data-maybe-add-ids! args))
 
 (defmethod sql.tx/create-table-sql :sparksql
   [driver {:keys [database-name]} {:keys [table-name field-definitions]}]
-  (let [quote-name    #(sql.u/quote-name driver :field (ddl.i/format-name driver %))
-        pk-field-name (quote-name (sql.tx/pk-field-name driver))]
-    (format "CREATE TABLE %s (%s %s, %s)"
+  (let [quote-name #(sql.u/quote-name driver :field (ddl.i/format-name driver %))]
+    (format "CREATE TABLE %s (%s)"
             (sql.tx/qualify-and-quote driver database-name table-name)
-            pk-field-name (sql.tx/pk-sql-type driver)
             (->> field-definitions
                  (map (fn [{:keys [field-name base-type]}]
                         (format "%s %s" (quote-name field-name) (if (map? base-type)

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

@@ -100,18 +100,25 @@
   "Dump a sequence of rows (as vectors) to a CSV file."
   [{:keys [field-definitions rows]} ^String filename]
   (try
-    (let [column-names (cons "id" (mapv :field-name field-definitions))
-          rows-with-id (for [[i row] (m/indexed rows)]
-                         (cons (inc i) (for [v row]
-                                         (value->csv v))))
-          csv-rows     (cons column-names rows-with-id)]
-      (try
-        (with-open [writer (java.io.FileWriter. (java.io.File. filename))]
-          (csv/write-csv writer csv-rows :quote? (constantly false)))
-        (catch Throwable e
-          (throw (ex-info "Error writing rows to CSV" {:rows (take 10 csv-rows)} e)))))
-    (catch Throwable e
-      (throw (ex-info "Error dumping rows to CSV" {:filename filename} e)))))
+   (let [has-custom-pk? (when-let [pk (sql.tx/fielddefs->pk-field-name field-definitions)]
+                          (not= "id" pk))
+         column-names   (cond->> (mapv :field-name field-definitions)
+                          (not has-custom-pk?)
+                          (cons "id"))
+         rows-with-id (for [[i row] (m/indexed rows)]
+                        (cond->> (for [v row]
+                                   (value->csv v))
+                          (not has-custom-pk?)
+                          (cons (inc i))))
+
+         csv-rows     (cons column-names rows-with-id)]
+     (try
+      (with-open [writer (java.io.FileWriter. (java.io.File. filename))]
+        (csv/write-csv writer csv-rows :quote? (constantly false)))
+      (catch Throwable e
+        (throw (ex-info "Error writing rows to CSV" {:rows (take 10 csv-rows)} e)))))
+   (catch Throwable e
+     (throw (ex-info "Error dumping rows to CSV" {:filename filename} e)))))
 
 (deftest dump-row-with-commas-to-csv-test
   (testing "Values with commas in them should get escaped correctly"

test/metabase/actions/test_util.clj

@@ -103,8 +103,15 @@
   `(do-with-dataset-definition actions-test-data (fn [] ~@body)))
 
 (defmacro with-temp-test-data
-  "Sets the current dataset to a freshly created dataset-definition that gets destroyed at the conclusion of `body`.
-   Use this to test destructive actions that may modify the data."
+  "Sets the current dataset to a freshly created table-definitions that gets destroyed at the conclusion of `body`.
+   Use this to test destructive actions that may modify the data.
+    (with-temp-test-data [[\"product\"
+                           [{:field-name \"name\" :base-type :type/Text}]
+                           [[\"Tesla Model S\"]]]
+                          [\"rating\"
+                           [{:field-name \"score\" :base-type :type/Integer}]
+                           [[5]]]]
+      ...)"
   {:style/indent :defn}
   [table-definitions & body]
   `(do-with-dataset-definition (apply tx/dataset-definition ~(str (gensym)) ~table-definitions) (fn [] ~@body)))

test/metabase/test/data/dataset_definition_test.clj

+(ns metabase.test.data.dataset-definition-test
+  (:require
+   [clojure.test :refer :all]
+   [metabase.driver :as driver]
+   [metabase.driver.ddl.interface :as ddl.i]
+   [metabase.test :as mt]
+   [toucan2.core :as t2]))
+
+(deftest dataset-with-custom-pk-test
+  (mt/test-drivers (disj (mt/sql-jdbc-drivers)
+                         ;; presto doesn't create PK for test data :) not sure why
+                         :presto-jdbc
+                         ;; creating db for athena is expensive and require some extra steps,
+                         ;; so it's not worth testing against, see [[metabase.test.data.athena/*allow-database-creation*]]
+                         :athena
+                         ;; there is no PK in sparksql
+                         :sparksql)
+    (mt/dataset (mt/dataset-definition "custom-pk"
+                  ["user"
+                   [{:field-name "custom_id" :base-type :type/Integer :pk? true}]
+                   [[1]]]
+                  ["group"
+                   [{:field-name "user_custom_id" :base-type :type/Integer :fk "user"}]
+                   [[1]]])
+      (let [user-fields  (t2/select [:model/Field :name :semantic_type :fk_target_field_id] :table_id (mt/id :user))
+            group-fields (t2/select [:model/Field :name :semantic_type :fk_target_field_id] :table_id (mt/id :group))
+            format-name  #(ddl.i/format-name driver/*driver* %)]
+        (testing "user.custom_id is a PK"
+          (is (= [{:name               (format-name "custom_id")
+                   :fk_target_field_id nil
+                   :semantic_type      :type/PK}]
+                 user-fields)))
+        (when-not (#{:sqlite} driver/*driver*) ;; our implement does not support adding fk for custom dataset yet
+          (testing "user_custom_id is a FK non user.custom_id"
+            (is (= #{{:name               (format-name "user_custom_id")
+                      :fk_target_field_id (mt/id :user :custom_id)
+                      :semantic_type      :type/FK}
+                     {:name               (format-name "id")
+                      :fk_target_field_id nil
+                      :semantic_type      :type/PK}}
+                   (set group-fields)))))))))

test/metabase/test/data/interface.clj

@@ -61,6 +61,7 @@
     ;; default is nullable
    [:not-null?         {:optional true} [:maybe :boolean]]
    [:unique?           {:optional true} [:maybe :boolean]]
+   [:pk?               {:optional true} [:maybe :boolean]]
    ;; should we create an index for this field?
    [:indexed?          {:optional true} [:maybe :boolean]]
    [:semantic-type     {:optional true} [:maybe ms/FieldSemanticOrRelationType]]

test/metabase/test/data/sql.clj

@@ -239,12 +239,16 @@
         inline-comment (inline-column-comment-sql driver field-comment)]
     (str/join " " (filter some? [field-name field-type not-null unique inline-comment]))))
 
+(defn fielddefs->pk-field-name
+  "Find the pk field name in fieldefs"
+  [fieldefs]
+  (->> fieldefs (filter :pk?) first :field-name))
+
 (defmethod create-table-sql :sql/test-extensions
   [driver {:keys [database-name], :as _dbdef} {:keys [table-name field-definitions table-comment]}]
-  (let [pk-field-name (format-and-quote-field-name driver (pk-field-name driver))]
-    (format "CREATE TABLE %s (%s %s, %s, PRIMARY KEY (%s)) %s;"
+  (let [pk-field-name (format-and-quote-field-name driver (fielddefs->pk-field-name field-definitions))]
+    (format "CREATE TABLE %s (%s, PRIMARY KEY (%s)) %s;"
             (qualify-and-quote driver database-name table-name)
-            pk-field-name (pk-sql-type driver)
             (str/join
              ", "
              (for [field-def field-definitions]
@@ -271,10 +275,20 @@
   tx/dispatch-on-driver-with-test-extensions
   :hierarchy #'driver/hierarchy)
 
+(defn- get-tabledef
+  [dbdef table-name]
+  (->> dbdef
+       :table-definitions
+       (filter #(= (:table-name %) table-name))
+       first))
+
 (defmethod add-fk-sql :sql/test-extensions
-  [driver {:keys [database-name]} {:keys [table-name]} {dest-table-name :fk, field-name :field-name}]
-  (let [quot            #(sql.u/quote-name driver %1 (ddl.i/format-name driver %2))
-        dest-table-name (name dest-table-name)]
+  [driver {:keys [database-name] :as dbdef} {:keys [table-name]} {dest-table-name :fk, field-name :field-name}]
+  (let [quot               #(sql.u/quote-name driver %1 (ddl.i/format-name driver %2))
+        dest-table-name    (name dest-table-name)
+        dest-table-pk-name (->> (get-tabledef dbdef dest-table-name)
+                                :field-definitions
+                                fielddefs->pk-field-name)]
     (format "ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s);"
             (qualify-and-quote driver database-name table-name)
             ;; limit FK constraint name to 30 chars since Oracle doesn't support names longer than that
@@ -285,7 +299,7 @@
               (quot :constraint fk-name))
             (quot :field field-name)
             (qualify-and-quote driver database-name dest-table-name)
-            (quot :field (pk-field-name driver)))))
+            (quot :field dest-table-pk-name))))
 
 (defmethod tx/count-with-template-tag-query :sql/test-extensions
   [driver table field _param-type]

test/metabase/test/data/sql/ddl.clj

@@ -39,12 +39,28 @@
   tx/dispatch-on-driver-with-test-extensions
   :hierarchy #'driver/hierarchy)
 
+(defn- add-pks-if-needed
+  "Add a pk for table that doesn't have one."
+  [driver tabledefs]
+  (map (fn [{:keys [field-definitions] :as tabledef}]
+         (if-not (some :pk? field-definitions)
+           (update tabledef :field-definitions #(cons (tx/map->FieldDefinition
+                                                       {:field-name    (sql.tx/pk-field-name driver)
+                                                        :base-type     {:native (sql.tx/pk-sql-type driver)}
+                                                        :semantic-type :type/PK
+                                                        :pk?           true})
+                                                      %))
+           tabledef))
+       tabledefs))
+
 (defmethod create-db-tables-ddl-statements :sql/test-extensions
-  [driver {:keys [table-definitions], :as dbdef} & _]
+  [driver dbdef & _]
   ;; Build combined statement for creating tables + FKs + comments
-  (let [statements (atom [])
-        add!       (fn [& stmnts]
-                     (swap! statements concat (filter some? stmnts)))]
+  (let [{:keys [table-definitions]
+         :as   dbdef}              (update dbdef :table-definitions (partial add-pks-if-needed driver))
+        statements                (atom [])
+        add!                      (fn [& stmnts]
+                                    (swap! statements concat (filter some? stmnts)))]
     (doseq [tabledef table-definitions]
       ;; Add the SQL for creating each Table
       (add! (sql.tx/drop-table-if-exists-sql driver dbdef tabledef)

test/metabase/test/data/sql_jdbc/load_data.clj

@@ -144,13 +144,23 @@
   "Implementation of `load-data!`. Insert rows in chunks of [[*chunk-size*]] (default 200) at a time."
   (make-load-data-fn load-data-chunked))
 
-(def ^{:arglists '([driver dbdef tabledef])} load-data-add-ids!
-  "Implementation of `load-data!`. Insert all rows at once; add IDs."
-  (make-load-data-fn load-data-add-ids))
-
-(def ^{:arglists '([driver dbdef tabledef])} load-data-add-ids-chunked!
-  "Implementation of `load-data!`. Insert rows in chunks of [[*chunk-size*]] (default 200) at a time; add IDs."
-  (make-load-data-fn load-data-add-ids load-data-chunked))
+(defn load-data-maybe-add-ids!
+  "Implementation of `load-data!`. Insert all rows at once;
+  Add IDs if tabledef does not contains PK."
+  [driver dbdef tabledef]
+  (let [load-data! (if-not (some :pk? (:field-definitions tabledef))
+                    (make-load-data-fn load-data-add-ids)
+                    (make-load-data-fn load-data-chunked))]
+    (load-data! driver dbdef tabledef)))
+
+(defn load-data-maybe-add-ids-chunked!
+  "Implementation of `load-data!`. Insert rows in chunks of [[*chunk-size*]] (default 200) at a time;
+  Add IDs if tabledef does not contains PK."
+  [driver dbdef tabledef]
+  (let [load-data! (if-not (some :pk? (:field-definitions tabledef))
+                    (make-load-data-fn load-data-add-ids load-data-chunked)
+                    (make-load-data-fn load-data-chunked))]
+    (load-data! driver dbdef tabledef)))
 
 ;; Default impl
 
@@ -197,7 +207,7 @@
 (defn create-db!
   "Default implementation of `create-db!` for SQL drivers."
   {:arglists '([driver dbdef & {:keys [skip-drop-db?]}])}
-  [driver {:keys [table-definitions], :as dbdef} & options]
+  [driver {:keys [table-definitions] :as dbdef} & options]
   ;; first execute statements to drop the DB if needed (this will do nothing if `skip-drop-db?` is true)
   (doseq [statement (apply ddl/drop-db-ddl-statements driver dbdef options)]
     (execute/execute-sql! driver :server dbdef statement))