Skip to content

Commit 06a48c4

Browse files
committed
Added pool example and tests for Psycopg 3 - closes #100
1 parent 25a3026 commit 06a48c4

File tree

3 files changed

+39
-1
lines changed

3 files changed

+39
-1
lines changed

README.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,15 @@ from pgvector.psycopg import register_vector
338338
register_vector(conn)
339339
```
340340

341+
For [connection pools](https://www.psycopg.org/psycopg3/docs/advanced/pool.html), use
342+
343+
```python
344+
def configure(conn):
345+
register_vector(conn)
346+
347+
pool = ConnectionPool(configure=configure)
348+
```
349+
341350
For [async connections](https://www.psycopg.org/psycopg3/docs/advanced/async.html), use
342351

343352
```python

requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ asyncpg
22
Django
33
numpy
44
peewee
5-
psycopg[binary]
5+
psycopg[binary,pool]
66
psycopg2-binary
77
pytest
88
pytest-asyncio

tests/test_psycopg.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import numpy as np
22
from pgvector.psycopg import register_vector, register_vector_async, Bit, HalfVector, SparseVector, Vector
33
import psycopg
4+
from psycopg_pool import ConnectionPool, AsyncConnectionPool
45
import pytest
56

67
conn = psycopg.connect(dbname='pgvector_python_test', autocommit=True)
@@ -176,6 +177,18 @@ def test_vector_array(self):
176177
assert np.array_equal(res[0][0], embeddings[0])
177178
assert np.array_equal(res[0][1], embeddings[1])
178179

180+
def test_pool(self):
181+
def configure(conn):
182+
register_vector(conn)
183+
184+
pool = ConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=True, configure=configure)
185+
186+
with pool.connection() as conn:
187+
res = conn.execute("SELECT '[1,2,3]'::vector").fetchone()
188+
assert np.array_equal(res[0], np.array([1, 2, 3]))
189+
190+
pool.close()
191+
179192
@pytest.mark.asyncio
180193
async def test_async(self):
181194
conn = await psycopg.AsyncConnection.connect(dbname='pgvector_python_test', autocommit=True)
@@ -195,3 +208,19 @@ async def test_async(self):
195208
assert np.array_equal(res[0][1], embedding)
196209
assert res[0][1].dtype == np.float32
197210
assert res[1][1] is None
211+
212+
@pytest.mark.asyncio
213+
async def test_async_pool(self):
214+
async def configure(conn):
215+
await register_vector_async(conn)
216+
217+
pool = AsyncConnectionPool(conninfo='postgres://localhost/pgvector_python_test', open=False, configure=configure)
218+
await pool.open()
219+
220+
async with pool.connection() as conn:
221+
async with conn.cursor() as cur:
222+
await cur.execute("SELECT '[1,2,3]'::vector")
223+
res = await cur.fetchone()
224+
assert np.array_equal(res[0], np.array([1, 2, 3]))
225+
226+
await pool.close()

0 commit comments

Comments
 (0)