Added tests for arrays with SQLAlchemy and Psycopg 3

pgvector · Feb 9, 2025 · 95403d5 · 95403d5
1 parent cfcc2ea
commit 95403d5
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
@@ -34,11 +34,21 @@
 
 
 @event.listens_for(psycopg2_array_engine, "connect")
-def connect(dbapi_connection, connection_record):
+def psycopg2_connect(dbapi_connection, connection_record):
     from pgvector.psycopg2 import register_vector
     register_vector(dbapi_connection, globally=False, arrays=True)
 
 
+if sqlalchemy_version > 1:
+    psycopg_array_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')
+    array_engines.append(psycopg_array_engine)
+
+    @event.listens_for(psycopg_array_engine, "connect")
+    def psycopg_connect(dbapi_connection, connection_record):
+        from pgvector.psycopg import register_vector
+        register_vector(dbapi_connection)
+
+
 Base = declarative_base()