Added schema option for asyncpg - closes #90

ankane · ankane · commit 3461181c17bb · 2024-09-20T19:20:00.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.3.4 (unreleased)
+
+- Added `schema` option for asyncpg
+
 ## 0.3.3 (2024-09-09)
 
 - Improved support for cursor factories with Psycopg 2
diff --git a/pgvector/asyncpg/register.py b/pgvector/asyncpg/register.py
@@ -1,9 +1,10 @@
 from ..utils import Vector, HalfVector, SparseVector
 
 
-async def register_vector(conn):
+async def register_vector(conn, schema='public'):
     await conn.set_type_codec(
         'vector',
+        schema=schema,
         encoder=Vector._to_db_binary,
         decoder=Vector._from_db_binary,
         format='binary'
@@ -12,13 +13,15 @@ async def register_vector(conn):
     try:
         await conn.set_type_codec(
             'halfvec',
+            schema=schema,
             encoder=HalfVector._to_db_binary,
             decoder=HalfVector._from_db_binary,
             format='binary'
         )
 
         await conn.set_type_codec(
             'sparsevec',
+            schema=schema,
             encoder=SparseVector._to_db_binary,
             decoder=SparseVector._from_db_binary,
             format='binary'
