Added pool example and test for asyncpg - closes pgvector#18

ankane · ankane · commit 34db96b3c1cb · 2023-06-28T18:30:12.000-04:00
diff --git a/README.md b/README.md
@@ -248,6 +248,15 @@ from pgvector.asyncpg import register_vector
 await register_vector(conn)
 ```
 
+or your pool
+
+```python
+async def init(conn):
+    await register_vector(conn)
+
+pool = await asyncpg.create_pool(..., init=register_vector)
+```
+
 Insert a vector
 
 ```python
diff --git a/tests/test_asyncpg.py b/tests/test_asyncpg.py
@@ -30,3 +30,25 @@ async def test_works(self):
         assert text_res[0]['embedding'] == '[1.5,2,3]'
 
         await conn.close()
+
+    @pytest.mark.asyncio
+    async def test_pool(self):
+        async def init(conn):
+            await register_vector(conn)
+
+        pool = await asyncpg.create_pool(database='pgvector_python_test', init=init)
+
+        async with pool.acquire() as conn:
+            await conn.execute('CREATE EXTENSION IF NOT EXISTS vector')
+            await conn.execute('DROP TABLE IF EXISTS item')
+            await conn.execute('CREATE TABLE item (id bigserial primary key, embedding vector(3))')
+
+            embedding = np.array([1.5, 2, 3])
+            await conn.execute("INSERT INTO item (embedding) VALUES ($1), (NULL)", embedding)
+
+            res = await conn.fetch("SELECT * FROM item ORDER BY id")
+            assert res[0]['id'] == 1
+            assert res[1]['id'] == 2
+            assert np.array_equal(res[0]['embedding'], embedding)
+            assert res[0]['embedding'].dtype == np.float32
+            assert res[1]['embedding'] is None
