diff --git a/src/diesel_ext.rs b/src/diesel_ext.rs index 3049387..628c91c 100644 --- a/src/diesel_ext.rs +++ b/src/diesel_ext.rs @@ -73,31 +73,33 @@ mod tests { table! { use diesel::sql_types::*; - items (id) { + diesel_items (id) { id -> Int4, embedding -> Nullable, } } - #[derive(PartialEq, Queryable)] - #[diesel(table_name = items)] + use diesel_items as items; + + #[derive(Queryable)] + #[diesel(table_name = diesel_items)] struct Item { pub id: i32, - pub embedding: Option, + pub embedding: Option, } #[derive(Insertable)] - #[diesel(table_name = items)] + #[diesel(table_name = diesel_items)] struct NewItem { - pub embedding: Option, + pub embedding: Option, } #[test] fn it_works() -> Result<(), diesel::result::Error> { let mut conn = PgConnection::establish("postgres://localhost/pgvector_rust_test").unwrap(); diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector").execute(&mut conn)?; - diesel::sql_query("DROP TABLE IF EXISTS items").execute(&mut conn)?; - diesel::sql_query("CREATE TABLE items (id serial PRIMARY KEY, embedding vector(3))") + diesel::sql_query("DROP TABLE IF EXISTS diesel_items").execute(&mut conn)?; + diesel::sql_query("CREATE TABLE diesel_items (id serial PRIMARY KEY, embedding vector(3))") .execute(&mut conn)?; let new_items = vec![ @@ -126,7 +128,11 @@ mod tests { .load::(&mut conn)?; assert_eq!( vec![1, 3, 2, 4], - neighbors.into_iter().map(|v| v.id).collect::>() + neighbors.iter().map(|v| v.id).collect::>() + ); + assert_eq!( + Some(Vector::from(vec![1.0, 1.0, 1.0])), + neighbors.first().unwrap().embedding ); let neighbors = items::table @@ -135,7 +141,7 @@ mod tests { .load::(&mut conn)?; assert_eq!( vec![2, 3, 1, 4], - neighbors.into_iter().map(|v| v.id).collect::>() + neighbors.iter().map(|v| v.id).collect::>() ); let neighbors = items::table @@ -144,7 +150,7 @@ mod tests { .load::(&mut conn)?; assert_eq!( vec![1, 2, 3, 4], - neighbors.into_iter().map(|v| v.id).collect::>() + neighbors.iter().map(|v| v.id).collect::>() ); let distances = items::table