Coverage for utilities/mptt_to_ltree.py: 56%
21 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 18:35 +0000
1"""
2Reusable SQL builders for migrating a django-mptt tree to a PostgreSQL ltree
3`path` (and optional `sort_path`) column.
5NetBox's core hierarchical models moved from django-mptt to ltree in v4.6. The
6per-table data backfill that migration performs is identical in shape for every
7tree, so the SQL is centralized here rather than copied into each app's
8migration. This also gives plugin maintainers a supported path for migrating
9their own MPTT models to `netbox.models.ltree.LtreeModel`; from a data migration:
11 from django.contrib.postgres.operations import CreateExtension
12 from django.db import migrations
13 from utilities.ltree import InstallLtreeTriggers
14 from utilities.mptt_to_ltree import assert_paths_populated_sql, populate_paths_sql
16 operations = [
17 CreateExtension('ltree'),
18 # ... AddField('path', nullable), [AddField('sort_path')], InstallLtreeTriggers(...) ...
19 migrations.RunSQL(
20 populate_paths_sql('myplugin_mymodel', sort_path=True),
21 reverse_sql=migrations.RunSQL.noop,
22 ),
23 migrations.RunSQL(
24 assert_paths_populated_sql('myplugin_mymodel'),
25 reverse_sql=migrations.RunSQL.noop,
26 ),
27 # ... AlterField('path' -> NOT NULL) ...
28 ]
30The values produced here must stay byte-identical to what the runtime triggers
31in `utilities.ltree` maintain: each path label is the row PK zero-padded to
32`_PATH_LABEL_WIDTH` chars, and `sort_path` is the chr(9) (TAB) separated chain of
33ancestor `name` values. Keep the two modules in sync if either changes.
34"""
36__all__ = (
37 'assert_paths_populated_sql',
38 'count_stale_rows_sql',
39 'populate_paths_sql',
40 'unreachable_rows_sql',
41)
43# Width to which each PK is zero-padded when used as an ltree label. Must match
44# the lpad() width used by the trigger functions in utilities.ltree (19 = max
45# bigint digit width) so that backfilled paths and trigger-maintained paths sort
46# and compare identically.
47_PATH_LABEL_WIDTH = 19
49# The SQL below names the ltree type and operators unqualified, so the extension's schema has to
50# be on the search_path. Put it there via set_config(..., true) — the function form of SET LOCAL —
51# rather than relying on the caller's path. Appending is a no-op when the schema is already on the
52# path, which also keeps repeated emissions idempotent.
53_ENSURE_LTREE_ON_PATH = """
54SELECT set_config('netbox.ltree_prior_search_path', current_setting('search_path'), true);
55SELECT set_config(
56 'search_path',
57 concat_ws(',', NULLIF(current_setting('search_path'), ''), quote_ident(n.nspname)),
58 true
59)
60FROM pg_extension e JOIN pg_namespace n ON n.oid = e.extnamespace
61WHERE e.extname = 'ltree' AND NOT n.nspname = ANY (current_schemas(true));
62"""
64# SET LOCAL persists to the end of the transaction, so restore the caller's value once the
65# ltree-dependent statements are done: a caller which narrows search_path to isolate unqualified
66# names must not have it left widened for the operations which follow.
67_RESTORE_SEARCH_PATH = """
68SELECT set_config('search_path', current_setting('netbox.ltree_prior_search_path'), true);
69"""
72def populate_paths_sql(table, *, sort_path=False):
73 """
74 Return SQL that backfills `path` (and `sort_path` when `sort_path=True`) for
75 every existing row in `table`, walking the tree from its roots
76 (parent_id IS NULL) downward via a single recursive CTE.
78 `path` is the chain of PK labels, each zero-padded to `_PATH_LABEL_WIDTH`
79 chars. `sort_path` is the chr(9) (TAB) separated chain of ancestor `name`
80 values, matching the `order_insertion_by=('name',)` semantics the triggers
81 maintain at runtime.
83 !!! warning
84 The UPDATE takes a row-exclusive lock on the entire table for the
85 duration of the statement. On large tables this can block writes for
86 minutes — plan a maintenance window accordingly.
88 !!! note
89 Run this inside a transaction, as an atomic migration does. The statements are
90 bracketed by set_config(..., true) — i.e. SET LOCAL — calls which put the ltree
91 extension's schema on the search_path and then restore the caller's value;
92 PostgreSQL discards SET LOCAL outside a transaction block.
93 """
94 if sort_path:
95 return _ENSURE_LTREE_ON_PATH + f"""
96WITH RECURSIVE t(id, parent_id, path, sort_path) AS (
97 SELECT id, parent_id,
98 lpad(id::text, {_PATH_LABEL_WIDTH}, '0')::ltree,
99 name::text
100 FROM "{table}" WHERE parent_id IS NULL
101 UNION ALL
102 SELECT r.id, r.parent_id,
103 t.path || lpad(r.id::text, {_PATH_LABEL_WIDTH}, '0')::ltree,
104 t.sort_path || chr(9) || r.name
105 FROM "{table}" r JOIN t ON r.parent_id = t.id
106)
107UPDATE "{table}" SET path = t.path, sort_path = t.sort_path
108FROM t WHERE "{table}".id = t.id;
109""" + _RESTORE_SEARCH_PATH
110 return _ENSURE_LTREE_ON_PATH + f"""
111WITH RECURSIVE t(id, parent_id, path) AS (
112 SELECT id, parent_id, lpad(id::text, {_PATH_LABEL_WIDTH}, '0')::ltree
113 FROM "{table}" WHERE parent_id IS NULL
114 UNION ALL
115 SELECT r.id, r.parent_id, t.path || lpad(r.id::text, {_PATH_LABEL_WIDTH}, '0')::ltree
116 FROM "{table}" r JOIN t ON r.parent_id = t.id
117)
118UPDATE "{table}" SET path = t.path FROM t WHERE "{table}".id = t.id;
119""" + _RESTORE_SEARCH_PATH
122def count_stale_rows_sql(table, sort_path=False):
123 """
124 Return SQL counting the rows in `table` whose `path` disagrees with the hierarchy, and
125 (when `sort_path` is set) the rows whose `sort_path` does.
127 A reparent leaves `path` wrong, a rename leaves `sort_path` wrong, and while the
128 cascade trigger is missing either can happen without the other, so both are counted
129 separately. Roots are checked against what `populate_paths_sql()` would give them (a
130 path of their own padded id, and a sort_path of their own name) and every other row
131 against its parent: a root has no parent to compare with, but it can still be wrong.
133 This answers "does this table need rebuilding", not "how many rows are damaged". Where
134 an object has moved, the objects below it agree with their own parent and are not
135 counted, though they are equally stale. Treat any non-zero result as the whole table
136 needing a rebuild, and do not use it to decide which rows to touch.
137 """
138 root_path = (
139 f'SELECT id FROM "{table}"'
140 f" WHERE parent_id IS NULL"
141 f" AND path <> lpad(id::text, {_PATH_LABEL_WIDTH}, '0')::ltree"
142 )
143 child_path = (
144 f'SELECT c.id FROM "{table}" c JOIN "{table}" p ON c.parent_id = p.id'
145 f" WHERE c.path <> p.path || lpad(c.id::text, {_PATH_LABEL_WIDTH}, '0')::ltree"
146 )
147 if sort_path:
148 root_sort_path = (
149 f'SELECT id FROM "{table}" WHERE parent_id IS NULL AND sort_path <> name'
150 )
151 child_sort_path = (
152 f'SELECT c.id FROM "{table}" c JOIN "{table}" p ON c.parent_id = p.id'
153 f' WHERE c.sort_path <> p.sort_path || chr(9) || c.name'
154 )
155 stale_sort_path = f'SELECT count(*) FROM ({root_sort_path} UNION ALL {child_sort_path}) s'
156 else:
157 stale_sort_path = 'SELECT 0'
159 return f"""
160SELECT
161 (SELECT count(*) FROM ({root_path} UNION ALL {child_path}) p) AS stale_paths,
162 ({stale_sort_path}) AS stale_sort_paths;
163"""
166def unreachable_rows_sql(table, limit):
167 """
168 Return SQL reporting the rows in `table` which no root can reach by following
169 `parent_id`: how many there are, and the first `limit` of their ids.
171 `populate_paths_sql()` seeds from `parent_id IS NULL` and walks downward, so it
172 rewrites only the rows reachable that way. Anything else it leaves untouched, which
173 makes an unreachable row an unrepaired one. Three shapes cause it: a cycle, a row
174 whose `parent_id` is its own id, and a `parent_id` referencing a row which does not
175 exist.
177 Callers which repair a populated table (rather than backfilling a fresh column, where
178 `assert_paths_populated_sql()` catches the same condition via the NULLs left behind)
179 should run this first and refuse if the count is non-zero: the parent relationships
180 have to be corrected before any path rebuild can produce a correct answer. The ids are
181 returned so that refusal can name rows to start from, rather than leaving the operator
182 to search the table for them.
183 """
184 return f"""
185WITH RECURSIVE reachable(id) AS (
186 SELECT id FROM "{table}" WHERE parent_id IS NULL
187 UNION ALL
188 SELECT c.id FROM "{table}" c JOIN reachable r ON c.parent_id = r.id
189)
190SELECT count(*), (array_agg(t.id ORDER BY t.id))[:{limit}]
191FROM "{table}" t
192WHERE NOT EXISTS (SELECT 1 FROM reachable r WHERE r.id = t.id);
193"""
196def assert_paths_populated_sql(table):
197 """
198 Return SQL that raises if any row in `table` still has a NULL `path` after
199 `populate_paths_sql()` runs.
201 The recursive CTE only reaches rows whose ancestry chains back to a
202 `parent_id IS NULL` root, so any row left with a NULL path points (directly
203 or transitively) at an orphan or cyclic parent_id. Catch that here, naming
204 the table and row count, rather than letting the subsequent
205 `AlterField(path -> NOT NULL)` abort opaquely inside ALTER COLUMN.
206 """
207 return f"""
208DO $$
209DECLARE missing bigint;
210BEGIN
211 SELECT count(*) INTO missing FROM "{table}" WHERE path IS NULL;
212 IF missing > 0 THEN
213 RAISE EXCEPTION
214 'ltree backfill left % rows in "{table}" with NULL path; '
215 'likely orphan parent_id references — resolve before re-running '
216 'this migration', missing;
217 END IF;
218END $$;
219"""