Skip to content

Commit 8afedf2

Browse files
[COST-8415] Add safe reconcile for stale partition tracking (#6402)
* [COST-8415] Add safe reconcile for stale partition tracking Add a dry-run management command that deletes partitioned_tables rows only when the physical relation is gone, and document the rare cleanup path for future partitioned-model removals. Co-authored-by: Cursor <cursoragent@cursor.com> * [COST-8415] Harden reconcile delete and reject empty --schema Use DELETE ... NOT EXISTS so absence is enforced in the same statement, reject an empty --schema value, and add type hints on the public API. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent 22ef8c5 commit 8afedf2

5 files changed

Lines changed: 382 additions & 0 deletions

File tree

‎docs/agent/partitioned-tables.md‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,3 +58,35 @@ number of large tenants have millions. An unindexed UPDATE/DELETE on
5858
take tens of seconds, blocking the API request under `@transaction.atomic`.
5959

6060
Always add an index on FK columns in partitioned tables.
61+
62+
## Stale `partitioned_tables` tracking (COST-8415)
63+
64+
Each tenant schema has a `partitioned_tables` registry. Creating/dropping partitions
65+
goes through that table and `trfn_partition_manager`. When a migration removes
66+
partitioned models (or drops parents) without cleaning the registry, tracking rows
67+
can remain. The monthly pre-create task **skips** those stale defaults (it does not
68+
delete them).
69+
70+
This is a **rare** event — typically when a provider or partitioned model family is
71+
removed from the product — not something that accumulates day-to-day.
72+
73+
### If you remove partitioned models
74+
75+
1. Drop/remove the physical tables (as the migration already does).
76+
2. Clear leftover registry rows **only** when the physical `schema.table_name` no
77+
longer exists. A `DELETE` on `partitioned_tables` while a same-named relation
78+
still exists will `DETACH` / `TRUNCATE` / `DROP` it via the trigger.
79+
3. Prefer the management command (dry-run first, then `--delete` after reviewing
80+
samples):
81+
82+
```bash
83+
# Inventory only (default)
84+
pipenv run python koku/manage.py reconcile_stale_partition_tracking --schema orgEXAMPLE
85+
# After reviewing log samples:
86+
pipenv run python koku/manage.py reconcile_stale_partition_tracking --schema orgEXAMPLE --delete
87+
```
88+
89+
Or call `koku.pg_partition.reconcile_stale_partition_tracking(schema_name, delete=...)`
90+
from a one-off ops path / migration RunPython with the same rule: never delete a
91+
row whose relation still exists in `pg_class` (name reuse / detached tables stay
92+
as report-only).

‎koku/koku/pg_partition.py‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import time
1414
import uuid
1515
from contextlib import contextmanager
16+
from typing import Any
1617

1718
import ciso8601
1819
from dateutil.relativedelta import relativedelta
@@ -2141,3 +2142,137 @@ def _handle_partitions(self, schema_name, table_names, start_date, end_date): #
21412142
LOG.debug(f"created = {created}")
21422143
if created:
21432144
LOG.info(f"Created partition {newpart.schema_name}.{newpart.table_name}")
2145+
2146+
2147+
# Classification labels for partition tracking reconciliation (COST-8415).
2148+
TRACKING_SAFE_ORPHAN = "safe_orphan"
2149+
TRACKING_PHYSICAL_PRESENT = "physical_present"
2150+
2151+
2152+
def relation_exists(schema_name: str, table_name: str) -> bool:
2153+
"""Return True if schema_name.table_name exists in pg_class."""
2154+
sql = """
2155+
SELECT EXISTS (
2156+
SELECT 1
2157+
FROM pg_class c
2158+
JOIN pg_namespace n
2159+
ON n.oid = c.relnamespace
2160+
WHERE n.nspname = %s::name
2161+
AND c.relname = %s::name
2162+
)
2163+
"""
2164+
cur = conn_execute(sql, (schema_name, table_name))
2165+
return bool(cur.fetchone()[0])
2166+
2167+
2168+
def _is_default_partition_parameters(partition_parameters: Any) -> bool:
2169+
if not isinstance(partition_parameters, dict):
2170+
return False
2171+
return bool(partition_parameters.get("default"))
2172+
2173+
2174+
def classify_partition_tracking(schema_name: str) -> list[dict[str, Any]]:
2175+
"""
2176+
Classify each partitioned_tables row in schema_name against PostgreSQL.
2177+
2178+
safe_orphan: tracking row whose physical table_name does not exist — DELETE is a
2179+
registry-only change (trfn_partition_manager takes no DDL action).
2180+
physical_present: relation exists (live, detached, or name reuse) — never auto-delete.
2181+
"""
2182+
with schema_context(schema_name):
2183+
records = list(PartitionedTable.objects.filter(schema_name=schema_name).order_by("table_name"))
2184+
2185+
classified = []
2186+
for record in records:
2187+
exists = relation_exists(schema_name, record.table_name)
2188+
classified.append(
2189+
{
2190+
"id": record.id,
2191+
"schema_name": schema_name,
2192+
"table_name": record.table_name,
2193+
"partition_of_table_name": record.partition_of_table_name,
2194+
"active": record.active,
2195+
"is_default": _is_default_partition_parameters(record.partition_parameters),
2196+
"classification": TRACKING_PHYSICAL_PRESENT if exists else TRACKING_SAFE_ORPHAN,
2197+
}
2198+
)
2199+
return classified
2200+
2201+
2202+
def _delete_safe_orphan_tracking(schema_name: str, candidate_ids: list[int]) -> int:
2203+
"""
2204+
Delete candidate tracking rows only when the physical relation is still absent.
2205+
2206+
Existence is enforced in the same DELETE statement (NOT EXISTS on pg_class) so a
2207+
concurrent CREATE cannot be selected for deletion after a separate Python check.
2208+
"""
2209+
if not candidate_ids:
2210+
return 0
2211+
sql = """
2212+
DELETE FROM partitioned_tables AS pt
2213+
WHERE pt.id = ANY(%s)
2214+
AND pt.schema_name = %s
2215+
AND NOT EXISTS (
2216+
SELECT 1
2217+
FROM pg_class c
2218+
JOIN pg_namespace n
2219+
ON n.oid = c.relnamespace
2220+
WHERE n.nspname = pt.schema_name::name
2221+
AND c.relname = pt.table_name::name
2222+
)
2223+
"""
2224+
with schema_context(schema_name):
2225+
cur = conn_execute(sql, (candidate_ids, schema_name))
2226+
return cur.rowcount
2227+
2228+
2229+
def reconcile_stale_partition_tracking(schema_name: str, *, delete: bool = False, sample: int = 5) -> dict[str, Any]:
2230+
"""
2231+
Inventory (and optionally delete) safe_orphan partition tracking rows for one schema.
2232+
2233+
Only deletes rows whose physical table is still absent at delete time. Never deletes
2234+
physical_present rows. Default is dry-run (delete=False).
2235+
2236+
Returns:
2237+
dict with keys: schema_name, delete, counts (by classification), deleted,
2238+
samples (up to sample rows per classification).
2239+
"""
2240+
classified = classify_partition_tracking(schema_name)
2241+
counts = {TRACKING_SAFE_ORPHAN: 0, TRACKING_PHYSICAL_PRESENT: 0}
2242+
samples: dict[str, list[dict[str, Any]]] = {TRACKING_SAFE_ORPHAN: [], TRACKING_PHYSICAL_PRESENT: []}
2243+
for row in classified:
2244+
label = row["classification"]
2245+
counts[label] = counts.get(label, 0) + 1
2246+
if len(samples.get(label, [])) < sample:
2247+
samples.setdefault(label, []).append(
2248+
{
2249+
"table_name": row["table_name"],
2250+
"partition_of_table_name": row["partition_of_table_name"],
2251+
"is_default": row["is_default"],
2252+
"active": row["active"],
2253+
}
2254+
)
2255+
2256+
deleted = 0
2257+
if delete:
2258+
safe_ids = [row["id"] for row in classified if row["classification"] == TRACKING_SAFE_ORPHAN]
2259+
deleted = _delete_safe_orphan_tracking(schema_name, safe_ids)
2260+
2261+
LOG.info(
2262+
log_json(
2263+
msg="partition tracking reconcile",
2264+
schema=schema_name,
2265+
delete=delete,
2266+
deleted=deleted,
2267+
safe_orphan=counts.get(TRACKING_SAFE_ORPHAN, 0),
2268+
physical_present=counts.get(TRACKING_PHYSICAL_PRESENT, 0),
2269+
samples=samples,
2270+
)
2271+
)
2272+
return {
2273+
"schema_name": schema_name,
2274+
"delete": delete,
2275+
"counts": counts,
2276+
"deleted": deleted,
2277+
"samples": samples,
2278+
}

‎koku/koku/test_pg_partition.py‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1489,3 +1489,115 @@ def hold_read_lock():
14891489
self.assertEqual(counts["created"], 1)
14901490
self.assertEqual(counts["missing"], [])
14911491
self.assertTrue(any("'outcome': 'created'" in line for line in logs.output))
1492+
1493+
1494+
class TestReconcileStalePartitionTracking(IamTestCase):
1495+
"""Safe cleanup of partitioned_tables rows whose physical tables are gone (COST-8415)."""
1496+
1497+
SCHEMA_NAME = "org1234567"
1498+
LIVE_TABLE = "__pg_partition_reconcile_live"
1499+
STALE_TABLE = "__pg_partition_reconcile_stale"
1500+
1501+
def setUp(self):
1502+
super().setUp()
1503+
self._create_partitioned_table(self.LIVE_TABLE)
1504+
self._create_partitioned_table(self.STALE_TABLE)
1505+
_execute(f"drop table {self.SCHEMA_NAME}.{self.STALE_TABLE} cascade")
1506+
1507+
def tearDown(self):
1508+
with schema_context(self.SCHEMA_NAME):
1509+
PartitionedTable.objects.filter(
1510+
schema_name=self.SCHEMA_NAME,
1511+
partition_of_table_name__in=[self.LIVE_TABLE, self.STALE_TABLE],
1512+
).delete()
1513+
_execute(f"drop table if exists {self.SCHEMA_NAME}.{self.LIVE_TABLE} cascade")
1514+
_execute(f"drop table if exists {self.SCHEMA_NAME}.{self.STALE_TABLE} cascade")
1515+
# Name-reuse helper tables from individual tests
1516+
_execute(f"drop table if exists {self.SCHEMA_NAME}.{self.STALE_TABLE}_default cascade")
1517+
super().tearDown()
1518+
1519+
def _create_partitioned_table(self, table_name):
1520+
_execute(
1521+
f"""
1522+
create table if not exists {self.SCHEMA_NAME}.{table_name} (
1523+
id bigserial,
1524+
usage_start date not null,
1525+
primary key (usage_start, id)
1526+
)
1527+
partition by range (usage_start);
1528+
"""
1529+
)
1530+
with schema_context(self.SCHEMA_NAME):
1531+
ppart.get_or_create_partition(
1532+
{
1533+
"schema_name": self.SCHEMA_NAME,
1534+
"table_name": f"{table_name}_default",
1535+
"partition_of_table_name": table_name,
1536+
"partition_type": PartitionedTable.RANGE,
1537+
"partition_col": "usage_start",
1538+
"partition_parameters": {"default": True},
1539+
"active": True,
1540+
}
1541+
)
1542+
1543+
def _tracking(self, table_name):
1544+
with schema_context(self.SCHEMA_NAME):
1545+
return PartitionedTable.objects.filter(schema_name=self.SCHEMA_NAME, table_name=table_name).first()
1546+
1547+
def test_dry_run_classifies_safe_orphan_and_does_not_delete(self):
1548+
result = ppart.reconcile_stale_partition_tracking(self.SCHEMA_NAME, delete=False)
1549+
self.assertFalse(result["delete"])
1550+
self.assertEqual(result["deleted"], 0)
1551+
self.assertGreaterEqual(result["counts"][ppart.TRACKING_SAFE_ORPHAN], 1)
1552+
self.assertIsNotNone(self._tracking(f"{self.STALE_TABLE}_default"))
1553+
self.assertIsNotNone(self._tracking(f"{self.LIVE_TABLE}_default"))
1554+
1555+
def test_delete_removes_only_safe_orphan_rows(self):
1556+
result = ppart.reconcile_stale_partition_tracking(self.SCHEMA_NAME, delete=True)
1557+
self.assertGreaterEqual(result["deleted"], 1)
1558+
self.assertIsNone(self._tracking(f"{self.STALE_TABLE}_default"))
1559+
self.assertIsNotNone(self._tracking(f"{self.LIVE_TABLE}_default"))
1560+
self.assertTrue(ppart.relation_exists(self.SCHEMA_NAME, f"{self.LIVE_TABLE}_default"))
1561+
1562+
def test_name_reuse_is_physical_present_and_not_deleted(self):
1563+
"""A same-named default attached to another parent must not be deleted."""
1564+
replacement = "__pg_partition_reconcile_replacement"
1565+
_execute(
1566+
f"""
1567+
create table {self.SCHEMA_NAME}.{replacement} (
1568+
id bigserial,
1569+
usage_start date not null,
1570+
primary key (usage_start, id)
1571+
)
1572+
partition by range (usage_start);
1573+
"""
1574+
)
1575+
_execute(
1576+
f"create table {self.SCHEMA_NAME}.{self.STALE_TABLE}_default "
1577+
f"partition of {self.SCHEMA_NAME}.{replacement} default"
1578+
)
1579+
try:
1580+
classified = ppart.classify_partition_tracking(self.SCHEMA_NAME)
1581+
stale_default = next(row for row in classified if row["table_name"] == f"{self.STALE_TABLE}_default")
1582+
self.assertEqual(stale_default["classification"], ppart.TRACKING_PHYSICAL_PRESENT)
1583+
self.assertEqual(stale_default["partition_of_table_name"], self.STALE_TABLE)
1584+
1585+
result = ppart.reconcile_stale_partition_tracking(self.SCHEMA_NAME, delete=True)
1586+
self.assertIsNotNone(self._tracking(f"{self.STALE_TABLE}_default"))
1587+
self.assertTrue(ppart.relation_exists(self.SCHEMA_NAME, f"{self.STALE_TABLE}_default"))
1588+
# Live default still present; name-reused stale default was not among deleted safe orphans
1589+
self.assertNotIn(
1590+
f"{self.STALE_TABLE}_default",
1591+
[s["table_name"] for s in result["samples"].get(ppart.TRACKING_SAFE_ORPHAN, [])],
1592+
)
1593+
finally:
1594+
_execute(f"drop table if exists {self.SCHEMA_NAME}.{replacement} cascade")
1595+
1596+
def test_delete_sql_skips_row_when_relation_exists(self):
1597+
"""DELETE ... NOT EXISTS must not remove tracking for a present relation."""
1598+
live_default = self._tracking(f"{self.LIVE_TABLE}_default")
1599+
self.assertIsNotNone(live_default)
1600+
deleted = ppart._delete_safe_orphan_tracking(self.SCHEMA_NAME, [live_default.id])
1601+
self.assertEqual(deleted, 0)
1602+
self.assertIsNotNone(self._tracking(f"{self.LIVE_TABLE}_default"))
1603+
self.assertTrue(ppart.relation_exists(self.SCHEMA_NAME, f"{self.LIVE_TABLE}_default"))
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
#
2+
# Copyright 2026 Red Hat Inc.
3+
# SPDX-License-Identifier: Apache-2.0
4+
#
5+
"""Inventory and safely delete orphaned partitioned_tables tracking rows (COST-8415).
6+
7+
Only deletes rows whose physical table does not exist in PostgreSQL. DELETE on
8+
partitioned_tables would otherwise DETACH/TRUNCATE/DROP a same-named relation via
9+
trfn_partition_manager — those rows are reported only.
10+
11+
Default is dry-run. Pass --delete after reviewing counts and samples.
12+
"""
13+
from __future__ import annotations
14+
15+
import argparse
16+
import logging
17+
from typing import Any
18+
19+
from django.core.management.base import BaseCommand
20+
from django.core.management.base import CommandError
21+
22+
from api.iam.models import Tenant
23+
from koku.pg_partition import reconcile_stale_partition_tracking
24+
from koku.pg_partition import TRACKING_PHYSICAL_PRESENT
25+
from koku.pg_partition import TRACKING_SAFE_ORPHAN
26+
27+
LOG = logging.getLogger(__name__)
28+
29+
30+
def _tenant_schema_names(*, schema: str | None = None) -> list[str]:
31+
excluded = {"public", Tenant._TEMPLATE_SCHEMA}
32+
if schema is not None:
33+
if not schema:
34+
raise CommandError("--schema must not be empty.")
35+
if schema in excluded:
36+
raise CommandError(f"Schema {schema} is not a customer tenant schema.")
37+
if not Tenant.objects.filter(schema_name=schema).exists():
38+
raise CommandError(f"Tenant schema {schema} not found.")
39+
return [schema]
40+
return [name for name in Tenant.objects.values_list("schema_name", flat=True) if name and name not in excluded]
41+
42+
43+
class Command(BaseCommand):
44+
help = (
45+
"Inventory (and optionally delete) partitioned_tables rows whose physical "
46+
"tables no longer exist. Dry-run by default; pass --delete to mutate."
47+
)
48+
49+
def add_arguments(self, parser: argparse.ArgumentParser) -> None:
50+
parser.add_argument(
51+
"--delete",
52+
action="store_true",
53+
default=False,
54+
help="Delete verified safe_orphan tracking rows. Without this flag, only report.",
55+
)
56+
parser.add_argument(
57+
"--schema",
58+
type=str,
59+
default=None,
60+
help="Limit to one tenant schema (default: all customer schemas).",
61+
)
62+
parser.add_argument(
63+
"--sample",
64+
type=int,
65+
default=5,
66+
help="Max sample rows to log per classification per schema (default: 5).",
67+
)
68+
69+
def handle(self, *args: Any, delete: bool, schema: str | None, sample: int, **kwargs: Any) -> None:
70+
if sample < 0:
71+
raise CommandError("--sample must be >= 0")
72+
73+
schemas = _tenant_schema_names(schema=schema)
74+
mode = "DELETE" if delete else "DRY RUN"
75+
self.stdout.write(f"partition tracking reconcile ({mode}) schemas={len(schemas)}")
76+
77+
totals = {TRACKING_SAFE_ORPHAN: 0, TRACKING_PHYSICAL_PRESENT: 0, "deleted": 0}
78+
for schema_name in schemas:
79+
result = reconcile_stale_partition_tracking(schema_name, delete=delete, sample=sample)
80+
totals[TRACKING_SAFE_ORPHAN] += result["counts"].get(TRACKING_SAFE_ORPHAN, 0)
81+
totals[TRACKING_PHYSICAL_PRESENT] += result["counts"].get(TRACKING_PHYSICAL_PRESENT, 0)
82+
totals["deleted"] += result["deleted"]
83+
84+
self.stdout.write(
85+
f"Done ({mode}): safe_orphan={totals[TRACKING_SAFE_ORPHAN]} "
86+
f"physical_present={totals[TRACKING_PHYSICAL_PRESENT]} deleted={totals['deleted']}"
87+
)
88+
if not delete and totals[TRACKING_SAFE_ORPHAN]:
89+
self.stdout.write("Review samples in logs, then re-run with --delete to remove safe_orphan rows only.")

0 commit comments

Comments
 (0)