|
| 1 | +"""rebuild inflated usage buckets |
| 2 | +
|
| 3 | +The aggregator summed amounts and durations separately and multiplied them |
| 4 | +afterwards, so both the JSONB ``resource_usage`` mirror and the normalized |
| 5 | +entries hold a cross product. Neither can be corrected in place, so both are |
| 6 | +rebuilt from ``kernel_usage_records``, which was never affected. |
| 7 | +
|
| 8 | +``usage_bucket_entries.amount`` becomes ``resource_usage``, matching what the |
| 9 | +JSONB mirror on ``kernel_usage_records`` and the three bucket tables already call |
| 10 | +this quantity, and drops its precision limit. It now holds the product directly |
| 11 | +rather than a factor readers had to multiply back out: a domain-level daily mem |
| 12 | +bucket runs past any fixed precision, and unconstrained NUMERIC has no ceiling. |
| 13 | +
|
| 14 | +``duration_seconds`` goes with it. It only ever existed to reconstitute that |
| 15 | +product, no reader consulted it on its own, and it counted kernel-seconds rather |
| 16 | +than wall-clock, so leaving it would leave a column that invites the same |
| 17 | +misreading the product form did. |
| 18 | +
|
| 19 | +Revision ID: c4a91d7e05b2 |
| 20 | +Revises: 5405ee0d8eed |
| 21 | +Create Date: 2026-07-20 00:00:00.000000 |
| 22 | +
|
| 23 | +""" |
| 24 | + |
| 25 | +import logging |
| 26 | +from datetime import date, timedelta |
| 27 | + |
| 28 | +import sqlalchemy as sa |
| 29 | +from alembic import op |
| 30 | + |
| 31 | +# revision identifiers, used by Alembic. |
| 32 | +revision = "c4a91d7e05b2" |
| 33 | +down_revision = "5405ee0d8eed" |
| 34 | +# Part of: NEXT_RELEASE_VERSION |
| 35 | +branch_labels = None |
| 36 | +depends_on = None |
| 37 | + |
| 38 | +log = logging.getLogger("alembic.runtime.migration") |
| 39 | + |
| 40 | + |
| 41 | +def upgrade() -> None: |
| 42 | + # Schema change: amount -> resource_usage (now the product, so no fixed |
| 43 | + # precision), and duration_seconds is no longer needed. |
| 44 | + op.alter_column( |
| 45 | + "usage_bucket_entries", |
| 46 | + "amount", |
| 47 | + new_column_name="resource_usage", |
| 48 | + existing_type=sa.Numeric(precision=24, scale=6), |
| 49 | + type_=sa.Numeric(), |
| 50 | + existing_nullable=False, |
| 51 | + ) |
| 52 | + op.drop_column("usage_bucket_entries", "duration_seconds") |
| 53 | + |
| 54 | + # Data change: rebuild the corrupted values from kernel_usage_records. |
| 55 | + # Delete every corrupted entry first (all three levels at once), then rebuild |
| 56 | + # each level. The three rebuilds share the same shape and differ only in which |
| 57 | + # table and key columns they use, so they are written out one by one below. |
| 58 | + conn = op.get_bind() |
| 59 | + window = _rebuildable_date_range(conn) |
| 60 | + if window is None: |
| 61 | + # No usage records to rebuild from (fresh install, or everything purged). |
| 62 | + return |
| 63 | + rebuild_from, rebuild_to = window |
| 64 | + _purge_corrupted_usage(conn, rebuild_from, rebuild_to) |
| 65 | + _rebuild_user_buckets(conn, rebuild_from, rebuild_to) |
| 66 | + _rebuild_project_buckets(conn, rebuild_from, rebuild_to) |
| 67 | + _rebuild_domain_buckets(conn, rebuild_from, rebuild_to) |
| 68 | + |
| 69 | + |
| 70 | +def downgrade() -> None: |
| 71 | + op.add_column( |
| 72 | + "usage_bucket_entries", |
| 73 | + sa.Column("duration_seconds", sa.Integer(), nullable=False, server_default="0"), |
| 74 | + ) |
| 75 | + # The rebuilt values are the correct ones and the inflated originals cannot |
| 76 | + # be reconstructed, so only the column definition is reverted. Values that |
| 77 | + # exceed the restored precision will fail the cast, which is the honest |
| 78 | + # outcome: they do not fit the old column. |
| 79 | + op.alter_column( |
| 80 | + "usage_bucket_entries", |
| 81 | + "resource_usage", |
| 82 | + new_column_name="amount", |
| 83 | + existing_type=sa.Numeric(), |
| 84 | + type_=sa.Numeric(precision=24, scale=6), |
| 85 | + existing_nullable=False, |
| 86 | + ) |
| 87 | + log.warning( |
| 88 | + "usage_bucket_entries is left corrupt by this downgrade: the column restored to " |
| 89 | + "'amount' now holds resource-seconds products, not raw amounts, and " |
| 90 | + "duration_seconds is reset to 0. Re-apply revision %s to rebuild the correct " |
| 91 | + "values from kernel_usage_records.", |
| 92 | + revision, |
| 93 | + ) |
| 94 | + |
| 95 | + |
| 96 | +def _rebuildable_date_range(conn: sa.engine.Connection) -> tuple[date, date] | None: |
| 97 | + """Return the date range that kernel_usage_records can faithfully rebuild. |
| 98 | +
|
| 99 | + The oldest retained day is excluded because retention purges by ``period_end`` |
| 100 | + and may have truncated it. Buckets outside the range keep their inflated |
| 101 | + values rather than being zeroed: they are unrecoverable, and zeroing them |
| 102 | + would destroy the only usage history left. |
| 103 | + """ |
| 104 | + row = conn.execute( |
| 105 | + sa.text( |
| 106 | + "SELECT min((period_start AT TIME ZONE 'UTC')::date) AS min_date, " |
| 107 | + " max((period_start AT TIME ZONE 'UTC')::date) AS max_date " |
| 108 | + "FROM kernel_usage_records" |
| 109 | + ) |
| 110 | + ).one() |
| 111 | + if row.min_date is None or row.max_date is None: |
| 112 | + return None |
| 113 | + rebuild_from = row.min_date + timedelta(days=1) |
| 114 | + if rebuild_from > row.max_date: |
| 115 | + return None |
| 116 | + return rebuild_from, row.max_date |
| 117 | + |
| 118 | + |
| 119 | +def _purge_corrupted_usage( |
| 120 | + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date |
| 121 | +) -> None: |
| 122 | + """Delete every corrupted entry the rebuild will replace, across all levels. |
| 123 | +
|
| 124 | + Each entry belongs to exactly one bucket via ``bucket_id``, so listing the |
| 125 | + in-window buckets of the three tables and deleting entries that point at them |
| 126 | + clears user, project and domain in a single statement. The rebuilds then |
| 127 | + start from a clean slate and only need to insert. |
| 128 | + """ |
| 129 | + conn.execute( |
| 130 | + sa.text( |
| 131 | + """ |
| 132 | + DELETE FROM usage_bucket_entries |
| 133 | + WHERE (bucket_id, bucket_type) IN ( |
| 134 | + SELECT id, 'user' FROM user_usage_buckets |
| 135 | + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to |
| 136 | + UNION ALL |
| 137 | + SELECT id, 'project' FROM project_usage_buckets |
| 138 | + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to |
| 139 | + UNION ALL |
| 140 | + SELECT id, 'domain' FROM domain_usage_buckets |
| 141 | + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to |
| 142 | + ) |
| 143 | + """ |
| 144 | + ), |
| 145 | + {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to}, |
| 146 | + ) |
| 147 | + |
| 148 | + |
| 149 | +def _rebuild_user_buckets(conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date) -> None: |
| 150 | + """Recompute user bucket entries and their JSONB mirror from kernel records.""" |
| 151 | + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} |
| 152 | + |
| 153 | + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds |
| 154 | + # that kernel_usage_records already stores correctly. capacity is refilled by |
| 155 | + # the next observation tick, so 0 here is safe. |
| 156 | + conn.execute( |
| 157 | + sa.text( |
| 158 | + """ |
| 159 | + INSERT INTO usage_bucket_entries |
| 160 | + (bucket_id, bucket_type, slot_name, resource_usage, capacity) |
| 161 | + SELECT user_usage_buckets.id, 'user', slot.key, SUM(slot.value::numeric), 0 |
| 162 | + FROM user_usage_buckets |
| 163 | + JOIN kernel_usage_records |
| 164 | + ON kernel_usage_records.user_uuid = user_usage_buckets.user_uuid |
| 165 | + AND kernel_usage_records.project_id = user_usage_buckets.project_id |
| 166 | + AND kernel_usage_records.resource_group_id = user_usage_buckets.resource_group_id |
| 167 | + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date |
| 168 | + = user_usage_buckets.period_start |
| 169 | + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot |
| 170 | + WHERE user_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 171 | + GROUP BY user_usage_buckets.id, slot.key |
| 172 | + """ |
| 173 | + ), |
| 174 | + params, |
| 175 | + ) |
| 176 | + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when |
| 177 | + # the bucket has no kernel records left to rebuild from. |
| 178 | + conn.execute( |
| 179 | + sa.text( |
| 180 | + """ |
| 181 | + UPDATE user_usage_buckets |
| 182 | + SET resource_usage = COALESCE( |
| 183 | + ( |
| 184 | + SELECT jsonb_object_agg( |
| 185 | + usage_bucket_entries.slot_name, |
| 186 | + usage_bucket_entries.resource_usage |
| 187 | + ) |
| 188 | + FROM usage_bucket_entries |
| 189 | + WHERE usage_bucket_entries.bucket_id = user_usage_buckets.id |
| 190 | + AND usage_bucket_entries.bucket_type = 'user' |
| 191 | + ), |
| 192 | + '{}'::jsonb |
| 193 | + ) |
| 194 | + WHERE user_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 195 | + """ |
| 196 | + ), |
| 197 | + params, |
| 198 | + ) |
| 199 | + |
| 200 | + |
| 201 | +def _rebuild_project_buckets( |
| 202 | + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date |
| 203 | +) -> None: |
| 204 | + """Recompute project bucket entries and their JSONB mirror from kernel records.""" |
| 205 | + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} |
| 206 | + |
| 207 | + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds |
| 208 | + # that kernel_usage_records already stores correctly. capacity is refilled by |
| 209 | + # the next observation tick, so 0 here is safe. |
| 210 | + conn.execute( |
| 211 | + sa.text( |
| 212 | + """ |
| 213 | + INSERT INTO usage_bucket_entries |
| 214 | + (bucket_id, bucket_type, slot_name, resource_usage, capacity) |
| 215 | + SELECT project_usage_buckets.id, 'project', slot.key, SUM(slot.value::numeric), 0 |
| 216 | + FROM project_usage_buckets |
| 217 | + JOIN kernel_usage_records |
| 218 | + ON kernel_usage_records.project_id = project_usage_buckets.project_id |
| 219 | + AND kernel_usage_records.resource_group_id = project_usage_buckets.resource_group_id |
| 220 | + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date |
| 221 | + = project_usage_buckets.period_start |
| 222 | + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot |
| 223 | + WHERE project_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 224 | + GROUP BY project_usage_buckets.id, slot.key |
| 225 | + """ |
| 226 | + ), |
| 227 | + params, |
| 228 | + ) |
| 229 | + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when |
| 230 | + # the bucket has no kernel records left to rebuild from. |
| 231 | + conn.execute( |
| 232 | + sa.text( |
| 233 | + """ |
| 234 | + UPDATE project_usage_buckets |
| 235 | + SET resource_usage = COALESCE( |
| 236 | + ( |
| 237 | + SELECT jsonb_object_agg( |
| 238 | + usage_bucket_entries.slot_name, |
| 239 | + usage_bucket_entries.resource_usage |
| 240 | + ) |
| 241 | + FROM usage_bucket_entries |
| 242 | + WHERE usage_bucket_entries.bucket_id = project_usage_buckets.id |
| 243 | + AND usage_bucket_entries.bucket_type = 'project' |
| 244 | + ), |
| 245 | + '{}'::jsonb |
| 246 | + ) |
| 247 | + WHERE project_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 248 | + """ |
| 249 | + ), |
| 250 | + params, |
| 251 | + ) |
| 252 | + |
| 253 | + |
| 254 | +def _rebuild_domain_buckets( |
| 255 | + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date |
| 256 | +) -> None: |
| 257 | + """Recompute domain bucket entries and their JSONB mirror from kernel records.""" |
| 258 | + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} |
| 259 | + |
| 260 | + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds |
| 261 | + # that kernel_usage_records already stores correctly. capacity is refilled by |
| 262 | + # the next observation tick, so 0 here is safe. |
| 263 | + conn.execute( |
| 264 | + sa.text( |
| 265 | + """ |
| 266 | + INSERT INTO usage_bucket_entries |
| 267 | + (bucket_id, bucket_type, slot_name, resource_usage, capacity) |
| 268 | + SELECT domain_usage_buckets.id, 'domain', slot.key, SUM(slot.value::numeric), 0 |
| 269 | + FROM domain_usage_buckets |
| 270 | + JOIN kernel_usage_records |
| 271 | + ON kernel_usage_records.domain_name = domain_usage_buckets.domain_name |
| 272 | + AND kernel_usage_records.resource_group_id = domain_usage_buckets.resource_group_id |
| 273 | + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date |
| 274 | + = domain_usage_buckets.period_start |
| 275 | + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot |
| 276 | + WHERE domain_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 277 | + GROUP BY domain_usage_buckets.id, slot.key |
| 278 | + """ |
| 279 | + ), |
| 280 | + params, |
| 281 | + ) |
| 282 | + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when |
| 283 | + # the bucket has no kernel records left to rebuild from. |
| 284 | + conn.execute( |
| 285 | + sa.text( |
| 286 | + """ |
| 287 | + UPDATE domain_usage_buckets |
| 288 | + SET resource_usage = COALESCE( |
| 289 | + ( |
| 290 | + SELECT jsonb_object_agg( |
| 291 | + usage_bucket_entries.slot_name, |
| 292 | + usage_bucket_entries.resource_usage |
| 293 | + ) |
| 294 | + FROM usage_bucket_entries |
| 295 | + WHERE usage_bucket_entries.bucket_id = domain_usage_buckets.id |
| 296 | + AND usage_bucket_entries.bucket_type = 'domain' |
| 297 | + ), |
| 298 | + '{}'::jsonb |
| 299 | + ) |
| 300 | + WHERE domain_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to |
| 301 | + """ |
| 302 | + ), |
| 303 | + params, |
| 304 | + ) |
0 commit comments