Added experimental support for psycopg3

ankane · ankane · commit cabb4e05bc5f · 2021-06-22T21:57:05.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.1.3 (unreleased)
 
 - Added support for asyncpg
+- Added experimental support for psycopg3
 
 ## 0.1.2 (2021-06-13)
 
diff --git a/README.md b/README.md
@@ -4,7 +4,7 @@
 
 Great for online recommendations :tada:
 
-Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [Psycopg 2](https://github.com/psycopg/psycopg2), and [asyncpg](https://github.com/MagicStack/asyncpg)
+Supports [Django](https://github.com/django/django), [SQLAlchemy](https://github.com/sqlalchemy/sqlalchemy), [Psycopg 2](https://github.com/psycopg/psycopg2), [Psycopg 3](https://github.com/psycopg/psycopg3), and [asyncpg](https://github.com/MagicStack/asyncpg)
 
 [![Build Status](https://github.com/ankane/pgvector-python/workflows/build/badge.svg?branch=master)](https://github.com/ankane/pgvector-python/actions)
 
@@ -21,6 +21,7 @@ And follow the instructions for your database library:
 - [Django](#django)
 - [SQLAlchemy](#sqlalchemy)
 - [Psycopg 2](#psycopg-2)
+- [Psycopg 3](#psycopg-3) [experimental]
 - [asyncpg](#asyncpg)
 
 Or check out some examples:
@@ -151,6 +152,29 @@ cur.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,))
 cur.fetchall()
 ```
 
+## psycopg3
+
+Register the vector type with your connection or cursor
+
+```python
+from pgvector.psycopg3 import register_vector
+
+register_vector(conn)
+```
+
+Insert a vector
+
+```python
+factors = np.array([1, 2, 3])
+cur.execute('INSERT INTO item (factors) VALUES (%s)', (factors,))
+```
+
+Get the nearest neighbors to a vector
+
+```python
+cur.execute('SELECT * FROM item ORDER BY factors <-> %s LIMIT 5', (factors,)).fetchall()
+```
+
 ## asyncpg
 
 Register the vector type with your connection
diff --git a/pgvector/psycopg3/__init__.py b/pgvector/psycopg3/__init__.py
@@ -0,0 +1,57 @@
+import numpy as np
+from psycopg3.adapt import Loader, Dumper
+from psycopg3.pq import Format
+from ..utils import from_db, from_db_binary, to_db, to_db_binary
+
+__all__ = ['register_vector']
+
+
+class VectorDumper(Dumper):
+
+    format = Format.TEXT
+
+    def dump(self, obj):
+        return to_db(obj).encode("utf8")
+
+
+class VectorBinaryDumper(VectorDumper):
+
+    format = Format.BINARY
+
+    def dump(self, obj):
+        return to_db_binary(obj)
+
+
+class VectorLoader(Loader):
+
+    format = Format.TEXT
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return from_db(data.decode("utf8"))
+
+
+class VectorBinaryLoader(VectorLoader):
+
+    format = Format.BINARY
+
+    def load(self, data):
+        if isinstance(data, memoryview):
+            data = bytes(data)
+        return from_db_binary(data)
+
+
+def register_vector(ctx):
+    cur = ctx.cursor() if hasattr(ctx, 'cursor') else ctx
+
+    try:
+        cur.execute('SELECT NULL::vector')
+        oid = cur.description[0][1]
+    except psycopg3.errors.UndefinedObject:
+        raise psycopg3.ProgrammingError('vector type not found in the database')
+
+    VectorDumper.register('numpy.ndarray', ctx)
+    VectorBinaryDumper.register('numpy.ndarray', ctx)
+    VectorLoader.register(oid, ctx)
+    VectorBinaryLoader.register(oid, ctx)
diff --git a/requirements.txt b/requirements.txt
@@ -5,3 +5,4 @@ psycopg2
 pytest
 pytest-asyncio
 SQLAlchemy
+git+https://github.com/psycopg/psycopg3.git#subdirectory=psycopg3
diff --git a/setup.py b/setup.py
@@ -12,6 +12,7 @@
         'pgvector.asyncpg',
         'pgvector.django',
         'pgvector.psycopg2',
+        'pgvector.psycopg3',
         'pgvector.sqlalchemy',
         'pgvector.utils'
     ],
diff --git a/tests/test_psycopg3.py b/tests/test_psycopg3.py
@@ -0,0 +1,35 @@
+import numpy as np
+from pgvector.psycopg3 import register_vector
+import psycopg3
+
+
+class TestPsycopg3(object):
+    def test_works(self):
+        conn = psycopg3.connect('dbname=pgvector_python_test')
+        conn.autocommit = True
+
+        cur = conn.cursor()
+        cur.execute('CREATE EXTENSION IF NOT EXISTS vector')
+        cur.execute('DROP TABLE IF EXISTS item')
+        cur.execute('CREATE TABLE item (id bigserial primary key, factors vector(3))')
+
+        register_vector(cur)
+
+        factors = np.array([1.5, 2, 3])
+        cur.execute("INSERT INTO item (factors) VALUES (%s), (NULL)", (factors,))
+
+        cur.execute("SELECT * FROM item ORDER BY id")
+        res = cur.fetchall()
+        assert res[0][0] == 1
+        assert res[1][0] == 2
+        assert np.array_equal(res[0][1], factors)
+        assert res[0][1].dtype == np.float32
+        assert res[1][1] is None
+
+        # binary format
+        binary_res = cur.execute("SELECT %b::vector", (factors,)).fetchone()[0]
+        assert np.array_equal(binary_res, factors)
+
+        # text format
+        text_res = cur.execute("SELECT %t::vector", (factors,)).fetchone()[0]
+        assert np.array_equal(text_res, factors)
