Closes #22973: fix ltree search_path (#22975)

netbox/utilities/mptt_to_ltree.py

@@ -44,6 +44,28 @@ __all__ = (
 # and compare identically.
 _PATH_LABEL_WIDTH = 19
 
+# The SQL below names the ltree type and operators unqualified, so the extension's schema has to
+# be on the search_path. Put it there via set_config(..., true) — the function form of SET LOCAL —
+# rather than relying on the caller's path. Appending is a no-op when the schema is already on the
+# path, which also keeps repeated emissions idempotent.
+_ENSURE_LTREE_ON_PATH = """
+SELECT set_config('netbox.ltree_prior_search_path', current_setting('search_path'), true);
+SELECT set_config(
+    'search_path',
+    concat_ws(',', NULLIF(current_setting('search_path'), ''), quote_ident(n.nspname)),
+    true
+)
+FROM pg_extension e JOIN pg_namespace n ON n.oid = e.extnamespace
+WHERE e.extname = 'ltree' AND NOT n.nspname = ANY (current_schemas(true));
+"""
+
+# SET LOCAL persists to the end of the transaction, so restore the caller's value once the
+# ltree-dependent statements are done: a caller which narrows search_path to isolate unqualified
+# names must not have it left widened for the operations which follow.
+_RESTORE_SEARCH_PATH = """
+SELECT set_config('search_path', current_setting('netbox.ltree_prior_search_path'), true);
+"""
+
 
 def populate_paths_sql(table, *, sort_path=False):
     """
@@ -60,9 +82,15 @@ def populate_paths_sql(table, *, sort_path=False):
         The UPDATE takes a row-exclusive lock on the entire table for the
         duration of the statement. On large tables this can block writes for
         minutes — plan a maintenance window accordingly.
+
+    !!! note
+        Run this inside a transaction, as an atomic migration does. The statements are
+        bracketed by set_config(..., true) — i.e. SET LOCAL — calls which put the ltree
+        extension's schema on the search_path and then restore the caller's value;
+        PostgreSQL discards SET LOCAL outside a transaction block.
     """
     if sort_path:
-        return f"""
+        return _ENSURE_LTREE_ON_PATH + f"""
 WITH RECURSIVE t(id, parent_id, path, sort_path) AS (
     SELECT id, parent_id,
            lpad(id::text, {_PATH_LABEL_WIDTH}, '0')::ltree,
@@ -76,8 +104,8 @@ WITH RECURSIVE t(id, parent_id, path, sort_path) AS (
 )
 UPDATE "{table}" SET path = t.path, sort_path = t.sort_path
 FROM t WHERE "{table}".id = t.id;
-"""
-    return f"""
+""" + _RESTORE_SEARCH_PATH
+    return _ENSURE_LTREE_ON_PATH + f"""
 WITH RECURSIVE t(id, parent_id, path) AS (
     SELECT id, parent_id, lpad(id::text, {_PATH_LABEL_WIDTH}, '0')::ltree
     FROM "{table}" WHERE parent_id IS NULL
@@ -86,7 +114,7 @@ WITH RECURSIVE t(id, parent_id, path) AS (
     FROM "{table}" r JOIN t ON r.parent_id = t.id
 )
 UPDATE "{table}" SET path = t.path FROM t WHERE "{table}".id = t.id;
-"""
+""" + _RESTORE_SEARCH_PATH
 
 
 def assert_paths_populated_sql(table):

netbox/utilities/tests/test_ltree.py

@@ -6,6 +6,7 @@ from django.test import TestCase
 from core.models import ObjectChange
 from dcim.models import Region, Site
 from tenancy.models import Contact, ContactGroup
+from utilities.mptt_to_ltree import populate_paths_sql
 
 
 def _path(*pks):
@@ -909,3 +910,52 @@ class NaturalSortSortPathTests(TestCase):
         )
         # Expected tree-flatten: parent, its child, then the unrelated root.
         self.assertEqual(names, ['nsP', 'nsPchild', 'nsP1'])
+
+
+class RestrictedSearchPathBackfillTests(TestCase):
+    """
+    The backfill SQL must resolve the ltree type itself, so that callers which apply
+    migrations one schema at a time (with the extension's schema off the search_path)
+    can still run it, and must leave the caller's search_path as it found it.
+    """
+
+    def _create_tree(self, cursor, table):
+        cursor.execute(
+            f'CREATE TABLE sp_test.{table} ('
+            f'id bigint PRIMARY KEY, parent_id bigint, path ltree, sort_path text, name text)'
+        )
+        cursor.execute(
+            f"INSERT INTO sp_test.{table} VALUES (1, NULL, NULL, NULL, 'root'), (2, 1, NULL, NULL, 'child')"
+        )
+
+    def test_backfill_with_extension_schema_off_search_path(self):
+        with connection.cursor() as cursor:
+            cursor.execute('CREATE SCHEMA sp_test')
+            cursor.execute('SET LOCAL search_path = sp_test, public')
+            self._create_tree(cursor, 'sorted')
+            self._create_tree(cursor, 'plain')
+
+            # As the migrations do: several tables in one statement batch, both variants,
+            # with the extension's schema absent from the path.
+            cursor.execute('SET LOCAL search_path = sp_test')
+            cursor.execute('\n'.join((
+                populate_paths_sql('sorted', sort_path=True),
+                populate_paths_sql('plain'),
+            )))
+            cursor.execute("SELECT current_setting('search_path')")
+            path_after = cursor.fetchone()[0]
+
+            cursor.execute('SET LOCAL search_path = sp_test, public')
+            cursor.execute('SELECT path::text, sort_path FROM sp_test.sorted ORDER BY id')
+            sorted_rows = cursor.fetchall()
+            cursor.execute('SELECT path::text FROM sp_test.plain ORDER BY id')
+            plain_rows = cursor.fetchall()
+
+        self.assertEqual(path_after, 'sp_test')
+
+        self.assertEqual(len(sorted_rows), 2)
+        self.assertEqual([r[0] for r in sorted_rows], [_path(1), _path(1, 2)])
+        self.assertEqual([r[1] for r in sorted_rows], ['root', 'root\tchild'])
+
+        self.assertEqual(len(plain_rows), 2)
+        self.assertEqual([r[0] for r in plain_rows], [_path(1), _path(1, 2)])