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.

1 files changed, 47 insertions, 0 deletions

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();
+    }
+}