Merge pull request diesel-rs#252 from weiznich/fix_211

Fix diesel-rs#211: Allow other primary keys than id for infer_schema

Author: sgrif

diesel_codegen/src/schema_inference/data_structures.rs

@@ -35,9 +35,9 @@ impl<ST> Queryable<ST, Pg> for ColumnInformation where
 #[cfg(feature = "sqlite")]
 impl<ST> Queryable<ST, Sqlite> for ColumnInformation where
     Sqlite: HasSqlType<ST>,
-    (i32, String, String, bool, Option<String>, i32): FromSqlRow<ST, Sqlite>,
+    (i32, String, String, bool, Option<String>, bool): FromSqlRow<ST, Sqlite>,
 {
-    type Row = (i32, String, String, bool, Option<String>, i32);
+    type Row = (i32, String, String, bool, Option<String>, bool);
 
     fn build(row: Self::Row) -> Self {
         ColumnInformation {

diesel_codegen/src/schema_inference/mod.rs

@@ -92,15 +92,30 @@ fn table_macro_call(
             Err(DummyResult::any(sp))
         }
         Ok(data) => {
-            let tokens = data.iter().map(|a| column_def_tokens(cx, a, &connection))
-                .collect::<Vec<_>>();
-            let table_name = str_to_ident(table_name);
-            let item = quote_item!(cx, table! {
-                $table_name {
-                    $tokens
+            let primary_keys = match get_primary_keys(connection, table_name) {
+                Ok(keys) => keys,
+                Err(_) =>{
+                    cx.span_err(sp, "error loading schema");
+                    return Err(DummyResult::any(sp));
                 }
-            }).unwrap();
-            Ok(item)
+            };
+            if primary_keys.len() != 1 {
+                cx.span_err(sp,
+                    &format!("table {} has {} primary keys, only one is currently supported",
+                        table_name, primary_keys.len()));
+                Err(DummyResult::any(sp))
+            } else {
+                let tokens = data.iter().map(|a| column_def_tokens(cx, a, &connection))
+                    .collect::<Vec<_>>();
+                let table_name = str_to_ident(table_name);
+                let primary_key = str_to_ident(&primary_keys[0]);
+                let item = quote_item!(cx, table! {
+                    $table_name ($primary_key) {
+                        $tokens
+                    }
+                }).unwrap();
+                Ok(item)
+            }
         }
     }
 }
@@ -201,3 +216,12 @@ fn determine_column_type(cx: &mut ExtCtxt, attr: &ColumnInformation, conn: &Infe
         InferConnection::Pg(_) => pg::determine_column_type(cx, attr),
     }
 }
+
+fn get_primary_keys(conn: &InferConnection, table_name: &str) -> QueryResult<Vec<String>> {
+    match *conn {
+        #[cfg(feature = "sqlite")]
+        InferConnection::Sqlite(ref c) => sqlite::get_primary_keys(c, table_name),
+        #[cfg(feature = "postgres")]
+        InferConnection::Pg(ref c) => pg::get_primary_keys(c, table_name),
+    }
+}

diesel_codegen/src/schema_inference/pg.rs

@@ -30,6 +30,15 @@ joinable!(pg_attribute -> pg_type (atttypid));
 select_column_workaround!(pg_attribute -> pg_type (attrelid, attname, atttypid, attnotnull, attnum, attisdropped));
 select_column_workaround!(pg_type -> pg_attribute (oid, typname));
 
+table! {
+    pg_index (indrelid) {
+        indrelid -> Oid,
+        indexrelid -> Oid,
+        indkey -> Array<SmallInt>,
+        indisprimary -> Bool,
+    }
+}
+
 table! {
     pg_class (oid) {
         oid -> Oid,
@@ -74,17 +83,32 @@ pub fn load_table_names(
 pub fn get_table_data(conn: &PgConnection, table_name: &str) -> QueryResult<Vec<ColumnInformation>> {
     use self::pg_attribute::dsl::*;
     use self::pg_type::dsl::{pg_type, typname};
-    let t_oid = try!(table_oid(conn, table_name));
+    use self::pg_class::dsl::*;
+
+    let table_oid = pg_class.select(oid).filter(relname.eq(table_name)).limit(1);
 
     pg_attribute.inner_join(pg_type)
         .select((attname, typname, attnotnull))
-        .filter(attrelid.eq(t_oid))
+        .filter(attrelid.eq_any(table_oid))
         .filter(attnum.gt(0).and(attisdropped.ne(true)))
         .order(attnum)
         .load(conn)
 }
 
-fn table_oid(conn: &PgConnection, table_name: &str) -> QueryResult<u32> {
+
+pub fn get_primary_keys(conn: &PgConnection, table_name: &str) -> QueryResult<Vec<String>> {
+    use self::pg_attribute::dsl::*;
+    use self::pg_index::dsl::{pg_index, indisprimary, indexrelid, indrelid};
     use self::pg_class::dsl::*;
-    pg_class.select(oid).filter(relname.eq(table_name)).first(conn)
+
+    let table_oid = pg_class.select(oid).filter(relname.eq(table_name)).limit(1);
+
+    let pk_query = pg_index.select(indexrelid)
+        .filter(indrelid.eq_any(table_oid))
+        .filter(indisprimary.eq(true));
+
+    pg_attribute.select(attname)
+        .filter(attrelid.eq_any(pk_query))
+        .order(attnum)
+        .load(conn)
 }

diesel_codegen/src/schema_inference/sqlite.rs

@@ -1,5 +1,6 @@
 use diesel::*;
-use diesel::sqlite::SqliteConnection;
+use diesel::sqlite::{SqliteConnection, Sqlite};
+use diesel::types::{HasSqlType, FromSqlRow};
 use syntax::ast;
 use syntax::codemap::Span;
 use syntax::ext::base::*;
@@ -14,7 +15,7 @@ table!{
         type_name -> VarChar,
         notnull -> Bool,
         dflt_value -> Nullable<VarChar>,
-        pk -> Integer,
+        pk -> Bool,
     }
 }
 
@@ -62,3 +63,38 @@ pub fn load_table_names(
         .filter(sql::<types::Bool>("type='table' AND name NOT LIKE '\\_\\_%'"));
     query.load(connection)
 }
+
+struct FullTableInfo {
+    _cid: i32,
+    name: String,
+    _type_name: String,
+    _not_null: bool,
+    _dflt_value: Option<String>,
+    primary_key: bool,
+}
+
+impl<ST> Queryable<ST, Sqlite> for FullTableInfo where
+    Sqlite: HasSqlType<ST>,
+    (i32, String, String, bool, Option<String>, bool): FromSqlRow<ST, Sqlite>,
+{
+    type Row = (i32, String, String, bool, Option<String>, bool);
+
+    fn build(row: Self::Row) -> Self {
+        FullTableInfo {
+            _cid: row.0,
+            name: row.1,
+            _type_name: row.2,
+            _not_null: row.3,
+            _dflt_value: row.4,
+            primary_key: row.5,
+        }
+    }
+}
+
+pub fn get_primary_keys(conn: &SqliteConnection, table_name: &str) -> QueryResult<Vec<String>> {
+    conn.execute_pragma::<pragma_table_info::SqlType, FullTableInfo>(
+        &format!("PRAGMA TABLE_INFO('{}')", table_name))
+        .map( |i| i.iter()
+               .filter_map(|i| if i.primary_key { Some(i.name.clone()) } else { None })
+               .collect())
+}