11import numpy as np
22from pgvector .psycopg import register_vector , register_vector_async , Bit , HalfVector , SparseVector , Vector
33import psycopg
4+ from psycopg_pool import ConnectionPool , AsyncConnectionPool
45import pytest
56
67conn = 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