From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Mon, 25 Sep 2023 22:26:56 +0000 (+0200)
Subject: fix: cache dynamic adapters created in register_array()
X-Git-Tag: 3.1.12~4^2~1
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=d7eea4989766ea92b0d7d51559ff11779ff3b872;p=thirdparty%2Fpsycopg.git

fix: cache dynamic adapters created in register_array()

If the base class is a C extension, the subclasses cannot be GC'd. This
results in a leak if register_array() is called in a loop.

Close #647
---

diff --git a/docs/news.rst b/docs/news.rst
index a88d9b3c4..ec9c44fa7 100644
--- a/docs/news.rst
+++ b/docs/news.rst
@@ -7,6 +7,13 @@
 ``psycopg`` release notes
 =========================
 
+Psycopg 3.1.12 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Fix memory leak when `~register_*()` functions are called repeatedly
+  (:ticket:`#647`).
+
+
 Current release
 ---------------
 
diff --git a/psycopg/psycopg/types/array.py b/psycopg/psycopg/types/array.py
index 419e3d41f..7a23a6333 100644
--- a/psycopg/psycopg/types/array.py
+++ b/psycopg/psycopg/types/array.py
@@ -311,41 +311,50 @@ def register_array(info: TypeInfo, context: Optional[AdaptContext] = None) -> No
     if not info.array_oid:
         raise ValueError(f"the type info {info} doesn't describe an array")
 
-    base: Type[Any]
     adapters = context.adapters if context else postgres.adapters
 
-    base = getattr(_psycopg, "ArrayLoader", ArrayLoader)
-    name = f"{info.name.title()}{base.__name__}"
-    attribs = {
-        "base_oid": info.oid,
-        "delimiter": info.delimiter.encode(),
-    }
-    loader = type(name, (base,), attribs)
+    loader = _make_loader(info.name, info.oid, info.delimiter)
     adapters.register_loader(info.array_oid, loader)
 
+    # No need to make a new loader because the binary datum has all the info.
     loader = getattr(_psycopg, "ArrayBinaryLoader", ArrayBinaryLoader)
     adapters.register_loader(info.array_oid, loader)
 
-    base = ListDumper
-    name = f"{info.name.title()}{base.__name__}"
-    attribs = {
-        "oid": info.array_oid,
-        "element_oid": info.oid,
-        "delimiter": info.delimiter.encode(),
-    }
-    dumper = type(name, (base,), attribs)
+    dumper = _make_dumper(info.name, info.oid, info.array_oid, info.delimiter)
     adapters.register_dumper(None, dumper)
 
-    base = ListBinaryDumper
-    name = f"{info.name.title()}{base.__name__}"
-    attribs = {
-        "oid": info.array_oid,
-        "element_oid": info.oid,
-    }
-    dumper = type(name, (base,), attribs)
+    dumper = _make_binary_dumper(info.name, info.oid, info.array_oid)
     adapters.register_dumper(None, dumper)
 
 
+# Cache all dynamically-generated types to avoid leaks in case the types
+# cannot be GC'd.
+
+
+@cache
+def _make_loader(name: str, oid: int, delimiter: str) -> Type[Loader]:
+    # Note: caching this function is really needed because, if the C extension
+    # is available, the resulting type cannot be GC'd, so calling
+    # register_array() in a loop results in a leak. See #647.
+    base = getattr(_psycopg, "ArrayLoader", ArrayLoader)
+    attribs = {"base_oid": oid, "delimiter": delimiter.encode()}
+    return type(f"{name.title()}{base.__name__}", (base,), attribs)
+
+
+@cache
+def _make_dumper(
+    name: str, oid: int, array_oid: int, delimiter: str
+) -> Type[BaseListDumper]:
+    attribs = {"oid": array_oid, "element_oid": oid, "delimiter": delimiter.encode()}
+    return type(f"{name.title()}ListDumper", (ListDumper,), attribs)
+
+
+@cache
+def _make_binary_dumper(name: str, oid: int, array_oid: int) -> Type[BaseListDumper]:
+    attribs = {"oid": array_oid, "element_oid": oid}
+    return type(f"{name.title()}ListBinaryDumper", (ListBinaryDumper,), attribs)
+
+
 def register_default_adapters(context: AdaptContext) -> None:
     # The text dumper is more flexible as it can handle lists of mixed type,
     # so register it later.
diff --git a/tests/types/test_array.py b/tests/types/test_array.py
index b0db15132..1f8c0f0de 100644
--- a/tests/types/test_array.py
+++ b/tests/types/test_array.py
@@ -1,3 +1,4 @@
+import gc
 from typing import List, Any
 from decimal import Decimal
 
@@ -10,6 +11,9 @@ from psycopg.adapt import PyFormat, Transformer, Dumper
 from psycopg.types import TypeInfo
 from psycopg._compat import prod
 from psycopg.postgres import types as builtins
+from psycopg.types.array import register_array
+
+from ..utils import gc_collect
 
 
 tests_str = [
@@ -335,3 +339,23 @@ def test_all_chars_with_bounds(conn, fmt_out):
     s = "".join(a)
     cur.execute("select '[0:1]={a,b}'::text[] || %s::text[]", ([s],))
     assert cur.fetchone()[0] == ["a", "b", s]
+
+
+def test_register_array_leak(conn):
+    info = TypeInfo.fetch(conn, "date")
+    ntypes = []
+    for i in range(2):
+        cur = conn.cursor()
+        register_array(info, cur)
+        cur.close()
+        del cur
+        gc_collect()
+
+        objs = gc.get_objects()
+        n = 0
+        for obj in objs:
+            if isinstance(obj, type):
+                n += 1
+        ntypes.append(n)
+
+    assert ntypes[0] == ntypes[1]