From 5b9fdb3b8917253242875ef30f6b13a064a224ae Mon Sep 17 00:00:00 2001
From: Teddy Wing
Date: Mon, 14 Mar 2022 21:14:57 +0100
Subject: Add `table_primary_key_column()`

Add a function to get the column name of the table's primary key. This
allows us to use a primary key value given by the user while saving
users the trouble of having to input the primary key column name.

We likely still want to offer the option of passing in a primary key
column name in case this function can't determine the table's primary
key, or if the table has a `UNIQUE` column we can query but no primary
key.
---
 src/sqlite.rs | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

(limited to 'src/sqlite.rs')

diff --git a/src/sqlite.rs b/src/sqlite.rs
index ba3c19d..c69bb81 100644
--- a/src/sqlite.rs
+++ b/src/sqlite.rs
@@ -69,3 +69,50 @@ pub fn get_column_names(
 
     Ok(column_names)
 }
+
+
+/// Get the name of the given table's primary key.
+pub fn table_primary_key_column(
+    dbconn: &rusqlite::Connection,
+    table_name: &str,
+) -> Result<String, crate::Error> {
+    let mut stmt = dbconn.prepare(r#"
+        SELECT "name"
+        FROM pragma_table_info(:table)
+        WHERE "pk" != 0;'
+    "#)?;
+
+    let pk_column: String = stmt.query_row(
+        &[(":table", table_name)],
+        |row| row.get(0),
+    )?;
+
+    Ok(pk_column)
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn table_primary_key_column_gets_primary_key_name() {
+        let conn = rusqlite::Connection::open_in_memory().unwrap();
+
+        conn.execute(
+            r#"
+                CREATE TABLE "test" (
+                    id INTEGER PRIMARY KEY,
+                    count INTEGER
+                );
+            "#,
+            []
+        ).unwrap();
+
+        let column_name = table_primary_key_column(&conn, "test").unwrap();
+
+        assert_eq!("id", column_name);
+
+        conn.close().unwrap();
+    }
+}
-- 
cgit v1.2.3