Fixed memory leak with register_vector and register_vector_async for Psycopg 3 - fixes #168

Author: ankane

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.5.2 (unreleased)
 
+- Fixed memory leak with `register_vector` and `register_vector_async` for Psycopg 3
 - Fixed error with `Bit` constructor and empty strings
 - Fixed equality for `HalfVector` with signed zeros
 

pgvector/psycopg/bit.py

@@ -1,3 +1,4 @@
+from functools import cache
 from psycopg import BaseConnection
 from psycopg.abc import Buffer
 from psycopg.adapt import Dumper
@@ -26,9 +27,19 @@ def register_bit_info(context: BaseConnection[Any], info: TypeInfo | None, /) ->
     info.register(context)
 
     # add oid to anonymous class for set_types
-    text_dumper = type('', (BitDumper,), {'oid': info.oid})
-    binary_dumper = type('', (BitBinaryDumper,), {'oid': info.oid})
+    text_dumper = bit_text_dumper(info.oid)
+    binary_dumper = bit_binary_dumper(info.oid)
 
     adapters = context.adapters
     adapters.register_dumper(Bit, text_dumper)
     adapters.register_dumper(Bit, binary_dumper)
+
+
+@cache
+def bit_text_dumper(oid: int):
+    return type('', (BitDumper,), {'oid': oid})
+
+
+@cache
+def bit_binary_dumper(oid: int):
+    return type('', (BitBinaryDumper,), {'oid': oid})

pgvector/psycopg/halfvec.py

@@ -1,3 +1,4 @@
+from functools import cache
 from psycopg import BaseConnection
 from psycopg.abc import Buffer
 from psycopg.adapt import Loader, Dumper
@@ -41,11 +42,21 @@ def register_halfvec_info(context: BaseConnection[Any], info: TypeInfo, /) -> No
     info.register(context)
 
     # add oid to anonymous class for set_types
-    text_dumper = type('', (HalfVectorDumper,), {'oid': info.oid})
-    binary_dumper = type('', (HalfVectorBinaryDumper,), {'oid': info.oid})
+    text_dumper = halfvec_text_dumper(info.oid)
+    binary_dumper = halfvec_binary_dumper(info.oid)
 
     adapters = context.adapters
     adapters.register_dumper(HalfVector, text_dumper)
     adapters.register_dumper(HalfVector, binary_dumper)
     adapters.register_loader(info.oid, HalfVectorLoader)
     adapters.register_loader(info.oid, HalfVectorBinaryLoader)
+
+
+@cache
+def halfvec_text_dumper(oid: int):
+    return type('', (HalfVectorDumper,), {'oid': oid})
+
+
+@cache
+def halfvec_binary_dumper(oid: int):
+    return type('', (HalfVectorBinaryDumper,), {'oid': oid})

pgvector/psycopg/sparsevec.py

@@ -1,3 +1,4 @@
+from functools import cache
 from psycopg import BaseConnection
 from psycopg.abc import Buffer
 from psycopg.adapt import Loader, Dumper
@@ -41,11 +42,21 @@ def register_sparsevec_info(context: BaseConnection[Any], info: TypeInfo, /) ->
     info.register(context)
 
     # add oid to anonymous class for set_types
-    text_dumper = type('', (SparseVectorDumper,), {'oid': info.oid})
-    binary_dumper = type('', (SparseVectorBinaryDumper,), {'oid': info.oid})
+    text_dumper = sparsevec_text_dumper(info.oid)
+    binary_dumper = sparsevec_binary_dumper(info.oid)
 
     adapters = context.adapters
     adapters.register_dumper(SparseVector, text_dumper)
     adapters.register_dumper(SparseVector, binary_dumper)
     adapters.register_loader(info.oid, SparseVectorLoader)
     adapters.register_loader(info.oid, SparseVectorBinaryLoader)
+
+
+@cache
+def sparsevec_text_dumper(oid: int):
+    return type('', (SparseVectorDumper,), {'oid': oid})
+
+
+@cache
+def sparsevec_binary_dumper(oid: int):
+    return type('', (SparseVectorBinaryDumper,), {'oid': oid})

pgvector/psycopg/vector.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from functools import cache
 import psycopg
 from psycopg import BaseConnection
 from psycopg.abc import Buffer
@@ -50,8 +51,8 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /)
     info.register(context)
 
     # add oid to anonymous class for set_types
-    text_dumper = type('', (VectorDumper,), {'oid': info.oid})
-    binary_dumper = type('', (VectorBinaryDumper,), {'oid': info.oid})
+    text_dumper = vector_text_dumper(info.oid)
+    binary_dumper = vector_binary_dumper(info.oid)
 
     adapters = context.adapters
     adapters.register_dumper(Vector, text_dumper)
@@ -60,3 +61,13 @@ def register_vector_info(context: BaseConnection[Any], info: TypeInfo | None, /)
     adapters.register_dumper('numpy.ndarray', binary_dumper)
     adapters.register_loader(info.oid, VectorLoader)
     adapters.register_loader(info.oid, VectorBinaryLoader)
+
+
+@cache
+def vector_text_dumper(oid: int):
+    return type('', (VectorDumper,), {'oid': oid})
+
+
+@cache
+def vector_binary_dumper(oid: int):
+    return type('', (VectorBinaryDumper,), {'oid': oid})