Added register_vector_async for psycopg3

ankane · akoshelev-fhl · ankane · commit 866f038acf22 · 2023-05-09T21:12:42.000-07:00
Co-authored-by: Alexander Koshelev &lt;a.koshelev@fhl.world&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## 0.1.7 (unreleased)
 
+- Added `register_vector_async` for psycopg3
 - Fixed `set_types` for psycopg3
 
 ## 0.1.6 (2022-05-22)
diff --git a/README.md b/README.md
@@ -180,6 +180,14 @@ from pgvector.psycopg import register_vector
 register_vector(conn)
 ```
 
+For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use [unreleased]
+
+```python
+from pgvector.psycopg import register_vector_async
+
+await register_vector_async(conn)
+```
+
 Insert a vector
 
 ```python
diff --git a/pgvector/psycopg/__init__.py b/pgvector/psycopg/__init__.py
@@ -45,6 +45,15 @@ def load(self, data):
 
 def register_vector(context):
     info = TypeInfo.fetch(context, 'vector')
+    register_vector_info(context, info)
+
+
+async def register_vector_async(context):
+    info = await TypeInfo.fetch(context, 'vector')
+    register_vector_info(context, info)
+
+
+def register_vector_info(context, info):
     if info is None:
         raise psycopg.ProgrammingError('vector type not found in the database')
     info.register(context)
diff --git a/tests/test_psycopg.py b/tests/test_psycopg.py
@@ -1,6 +1,7 @@
 import numpy as np
-from pgvector.psycopg import register_vector
+from pgvector.psycopg import register_vector, register_vector_async
 import psycopg
+import pytest
 
 conn = psycopg.connect(dbname='pgvector_python_test')
 conn.autocommit = True
@@ -70,3 +71,23 @@ def test_binary_copy_set_types(self):
         with cur.copy("COPY item (id, embedding) FROM STDIN WITH (FORMAT BINARY)") as copy:
             copy.set_types(['int8', 'vector'])
             copy.write_row([1, embedding])
+
+    @pytest.mark.asyncio
+    async def test_async(self):
+        conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True)
+
+        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))')
+
+        await register_vector_async(conn)
+
+        embedding = np.array([1.5, 2, 3])
+        await conn.execute('INSERT INTO item (embedding) VALUES (%s), (NULL)', (embedding,))
+
+        async with conn.cursor() as cur:
+            await cur.execute('SELECT * FROM item ORDER BY id')
+            res = await cur.fetchall()
+            assert np.array_equal(res[0][1], embedding)
+            assert res[0][1].dtype == np.float32
+            assert res[1][1] is None
