Skip to content

Commit bbbf7b3

Browse files
committed
Added test for tokio-postgres - #23
1 parent bab73fc commit bbbf7b3

File tree

2 files changed

+38
-0
lines changed

2 files changed

+38
-0
lines changed

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ half = { version = "2", default-features = false, optional = true }
2222

2323
[dev-dependencies]
2424
postgres = { version = "0.19", default-features = false }
25+
tokio-postgres = { version = "0.7", default-features = false }
2526
diesel = { version = "2", default-features = false, features = ["32-column-tables"] }
2627
sqlx = { version = "0", default-features = false, features = ["runtime-tokio", "tls-native-tls"] }
2728
tokio = { version = "1", features = ["full"] }

src/postgres_ext/vector.rs

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,4 +121,41 @@ mod tests {
121121
row.get("schema"),
122122
))
123123
}
124+
125+
#[tokio::test]
126+
async fn tokio_works() -> Result<(), tokio_postgres::Error> {
127+
let (client, connection) =
128+
tokio_postgres::connect("host=localhost dbname=pgvector_rust_test", tokio_postgres::NoTls).await?;
129+
130+
tokio::spawn(async move {
131+
if let Err(e) = connection.await {
132+
eprintln!("connection error: {}", e);
133+
}
134+
});
135+
136+
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[]).await?;
137+
client.execute("DROP TABLE IF EXISTS tokio_postgres_items", &[]).await?;
138+
client.execute(
139+
"CREATE TABLE tokio_postgres_items (id bigserial PRIMARY KEY, embedding vector(3))",
140+
&[],
141+
).await?;
142+
143+
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
144+
let vec2 = Vector::from(vec![4.0, 5.0, 6.0]);
145+
client.execute(
146+
"INSERT INTO tokio_postgres_items (embedding) VALUES ($1), ($2), (NULL)",
147+
&[&vec, &vec2],
148+
).await?;
149+
150+
let query_vec = Vector::from(vec![3.0, 1.0, 2.0]);
151+
let row = client.query_one(
152+
"SELECT embedding FROM tokio_postgres_items ORDER BY embedding <-> $1 LIMIT 1",
153+
&[&query_vec],
154+
).await?;
155+
let res_vec: Vector = row.get(0);
156+
assert_eq!(vec, res_vec);
157+
assert_eq!(vec![1.0, 2.0, 3.0], res_vec.to_vec());
158+
159+
Ok(())
160+
}
124161
}

0 commit comments

Comments
 (0)