Skip to content

Commit 1bb4810

Browse files
gm-e2bgm-e2b
committed
[EN-1038] feat(db): split template disk entitlements
Add nullable tier entitlements for the default template free-space target and maximum logical rootfs size, plus an add-on contribution to the maximum. Validate non-negative defaults and add-on increments, positive maxima, and that each default does not exceed its maximum. Backfill existing add-ons so their maximum contribution initially matches extra_disk_mb. Keep the new field nullable and fall back to extra_disk_mb for writes from legacy services during rollout. Leave tier population to the audited follow-up data migration. Extend team_limits in place, preserving its existing columns and appending the effective default_free_disk_size_mb and max_disk_size_mb values computed from active add-ons. Regenerate sqlc queries and models with nullable int64 entitlement fields, and cover the schema shape, security-invoker setting, column compatibility, and tier-plus-add-on calculations in the migration test. Part of: EN-1038
1 parent 73f4eb5 commit 1bb4810

9 files changed

Lines changed: 265 additions & 5 deletions

File tree

‎packages/db/migrations/20260714091414_add_template_disk_entitlements.sql‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
-- +goose Up
2+
-- +goose StatementBegin
3+
4+
ALTER TABLE "public"."tiers"
5+
-- Free-rootfs target in the existing MiB convention.
6+
ADD COLUMN IF NOT EXISTS "default_free_disk_size_mb" bigint,
7+
-- Total logical-rootfs ceiling before active add-ons.
8+
ADD COLUMN IF NOT EXISTS "max_disk_size_mb" bigint;
9+
10+
ALTER TABLE "public"."addons"
11+
-- Total-ceiling increment in the existing MiB convention.
12+
ADD COLUMN IF NOT EXISTS "extra_max_disk_size_mb" bigint;
13+
14+
-- Populate existing tiers in every environment before the effective limits are exposed.
15+
UPDATE "public"."tiers"
16+
SET "default_free_disk_size_mb" = "disk_mb",
17+
"max_disk_size_mb" = "disk_mb" + 25000;
18+
19+
-- Populate rows that predate the schema expansion in every environment. Keep the column nullable
20+
-- because legacy writers may still omit it during the rollout.
21+
UPDATE "public"."addons"
22+
SET "extra_max_disk_size_mb" = "extra_disk_mb"
23+
WHERE "extra_max_disk_size_mb" IS NULL;
24+
25+
ALTER TABLE "public"."tiers"
26+
ADD CONSTRAINT "tiers_default_free_disk_size_mb_check"
27+
CHECK (default_free_disk_size_mb >= 0),
28+
ADD CONSTRAINT "tiers_max_disk_size_mb_check"
29+
CHECK (max_disk_size_mb > 0),
30+
ADD CONSTRAINT "tiers_default_free_disk_size_lte_max_check"
31+
CHECK (default_free_disk_size_mb <= max_disk_size_mb);
32+
33+
ALTER TABLE "public"."addons"
34+
ADD CONSTRAINT "addons_extra_max_disk_size_mb_check"
35+
CHECK (extra_max_disk_size_mb >= 0);
36+
37+
-- +goose StatementEnd
38+
39+
-- +goose Down
40+
-- +goose StatementBegin
41+
42+
ALTER TABLE "public"."tiers"
43+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_lte_max_check",
44+
DROP CONSTRAINT IF EXISTS "tiers_max_disk_size_mb_check",
45+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_mb_check",
46+
DROP COLUMN IF EXISTS "max_disk_size_mb",
47+
DROP COLUMN IF EXISTS "default_free_disk_size_mb";
48+
49+
ALTER TABLE "public"."addons"
50+
DROP CONSTRAINT IF EXISTS "addons_extra_max_disk_size_mb_check",
51+
DROP COLUMN IF EXISTS "extra_max_disk_size_mb";
52+
53+
-- +goose StatementEnd

‎packages/db/migrations/20260714091415_extend_team_limits.sql‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
-- +goose Up
2+
-- +goose StatementBegin
3+
4+
CREATE OR REPLACE VIEW "public"."team_limits"
5+
WITH (security_invoker=on) AS
6+
SELECT
7+
t.id,
8+
tier.max_length_hours,
9+
(tier.concurrent_instances + a.extra_concurrent_sandboxes) AS concurrent_sandboxes,
10+
(tier.concurrent_template_builds + a.extra_concurrent_template_builds) AS concurrent_template_builds,
11+
(tier.max_vcpu + a.extra_max_vcpu) AS max_vcpu,
12+
(tier.max_ram_mb + a.extra_max_ram_mb) AS max_ram_mb,
13+
(tier.disk_mb + a.extra_disk_mb) AS disk_mb,
14+
(tier.events_ttl_days + a.extra_events_ttl_days) AS events_ttl_days,
15+
(tier.default_free_disk_size_mb + a.extra_disk_mb)::bigint AS default_free_disk_size_mb,
16+
(tier.max_disk_size_mb + a.extra_max_disk_size_mb)::bigint AS max_disk_size_mb
17+
FROM "public"."teams" t
18+
JOIN "public"."tiers" tier ON t.tier = tier.id
19+
LEFT JOIN LATERAL (
20+
SELECT COALESCE(SUM(extra_concurrent_sandboxes), 0)::bigint AS extra_concurrent_sandboxes,
21+
COALESCE(SUM(extra_concurrent_template_builds), 0)::bigint AS extra_concurrent_template_builds,
22+
COALESCE(SUM(extra_max_vcpu), 0)::bigint AS extra_max_vcpu,
23+
COALESCE(SUM(extra_max_ram_mb), 0)::bigint AS extra_max_ram_mb,
24+
COALESCE(SUM(extra_disk_mb), 0)::bigint AS extra_disk_mb,
25+
COALESCE(SUM(extra_events_ttl_days), 0)::bigint AS extra_events_ttl_days,
26+
COALESCE(SUM(COALESCE(extra_max_disk_size_mb, extra_disk_mb)), 0)::bigint AS extra_max_disk_size_mb
27+
FROM "public"."addons" addon
28+
WHERE addon.team_id = t.id
29+
AND addon.valid_from <= now()
30+
AND (addon.valid_to IS NULL OR addon.valid_to > now())
31+
) a ON true;
32+
33+
-- +goose StatementEnd
34+
35+
-- +goose Down
36+
-- +goose StatementBegin
37+
38+
-- PostgreSQL cannot remove appended view columns with CREATE OR REPLACE, so restore the previous
39+
-- definition transactionally. Concurrent readers cannot observe the view between these statements.
40+
DROP VIEW IF EXISTS "public"."team_limits";
41+
42+
CREATE VIEW "public"."team_limits"
43+
WITH (security_invoker=on) AS
44+
SELECT
45+
t.id,
46+
tier.max_length_hours,
47+
(tier.concurrent_instances + a.extra_concurrent_sandboxes) AS concurrent_sandboxes,
48+
(tier.concurrent_template_builds + a.extra_concurrent_template_builds) AS concurrent_template_builds,
49+
(tier.max_vcpu + a.extra_max_vcpu) AS max_vcpu,
50+
(tier.max_ram_mb + a.extra_max_ram_mb) AS max_ram_mb,
51+
(tier.disk_mb + a.extra_disk_mb) AS disk_mb,
52+
(tier.events_ttl_days + a.extra_events_ttl_days) AS events_ttl_days
53+
FROM "public"."teams" t
54+
JOIN "public"."tiers" tier ON t.tier = tier.id
55+
LEFT JOIN LATERAL (
56+
SELECT COALESCE(SUM(extra_concurrent_sandboxes), 0)::bigint AS extra_concurrent_sandboxes,
57+
COALESCE(SUM(extra_concurrent_template_builds), 0)::bigint AS extra_concurrent_template_builds,
58+
COALESCE(SUM(extra_max_vcpu), 0)::bigint AS extra_max_vcpu,
59+
COALESCE(SUM(extra_max_ram_mb), 0)::bigint AS extra_max_ram_mb,
60+
COALESCE(SUM(extra_disk_mb), 0)::bigint AS extra_disk_mb,
61+
COALESCE(SUM(extra_events_ttl_days), 0)::bigint AS extra_events_ttl_days
62+
FROM "public"."addons" addon
63+
WHERE addon.team_id = t.id
64+
AND addon.valid_from <= now()
65+
AND (addon.valid_to IS NULL OR addon.valid_to > now())
66+
) a ON true;
67+
68+
-- +goose StatementEnd

‎packages/db/pkg/auth/queries/get_team.sql.go‎

Lines changed: 12 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/pkg/auth/queries/models.go‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/pkg/dashboard/queries/get_dashboard_teams_with_users_teams_with_tier.sql.go‎

Lines changed: 3 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/pkg/dashboard/queries/models.go‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/pkg/tests/disk_entitlements_migration_test.go‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
package tests
2+
3+
import (
4+
"database/sql"
5+
"testing"
6+
7+
"github.com/google/uuid"
8+
"github.com/stretchr/testify/require"
9+
10+
"github.com/e2b-dev/infra/packages/db/pkg/testutils"
11+
testqueries "github.com/e2b-dev/infra/packages/db/pkg/testutils/queries"
12+
)
13+
14+
func TestDiskEntitlementsMigration(t *testing.T) {
15+
t.Parallel()
16+
17+
db := testutils.SetupDatabase(t)
18+
ctx := t.Context()
19+
20+
sqlDB, err := sql.Open("pgx", db.ConnStr())
21+
require.NoError(t, err)
22+
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
23+
24+
var columnCount int64
25+
var nullableWithoutDefaults bool
26+
err = sqlDB.QueryRowContext(ctx, `
27+
SELECT COUNT(*), BOOL_AND(data_type = 'bigint' AND is_nullable = 'YES' AND column_default IS NULL)
28+
FROM information_schema.columns
29+
WHERE table_schema = 'public'
30+
AND (table_name, column_name) IN (
31+
('tiers', 'default_free_disk_size_mb'),
32+
('tiers', 'max_disk_size_mb'),
33+
('addons', 'extra_max_disk_size_mb')
34+
)
35+
`).Scan(&columnCount, &nullableWithoutDefaults)
36+
require.NoError(t, err)
37+
require.Equal(t, int64(3), columnCount)
38+
require.True(t, nullableWithoutDefaults)
39+
40+
var tierCount, invalidTierCount int64
41+
err = sqlDB.QueryRowContext(ctx, `
42+
SELECT COUNT(*), COUNT(*) FILTER (
43+
WHERE default_free_disk_size_mb IS DISTINCT FROM disk_mb
44+
OR max_disk_size_mb IS DISTINCT FROM disk_mb + 25000
45+
)
46+
FROM public.tiers
47+
`).Scan(&tierCount, &invalidTierCount)
48+
require.NoError(t, err)
49+
require.NotZero(t, tierCount)
50+
require.Zero(t, invalidTierCount)
51+
52+
var viewColumns string
53+
err = sqlDB.QueryRowContext(ctx, `
54+
SELECT string_agg(column_name, ',' ORDER BY ordinal_position)
55+
FROM information_schema.columns
56+
WHERE table_schema = 'public' AND table_name = 'team_limits'
57+
`).Scan(&viewColumns)
58+
require.NoError(t, err)
59+
60+
const legacyColumns = "id,max_length_hours,concurrent_sandboxes,concurrent_template_builds," +
61+
"max_vcpu,max_ram_mb,disk_mb,events_ttl_days"
62+
require.Equal(t, legacyColumns+",default_free_disk_size_mb,max_disk_size_mb", viewColumns)
63+
64+
var securityInvoker bool
65+
err = sqlDB.QueryRowContext(ctx, `
66+
SELECT COALESCE(reloptions @> ARRAY['security_invoker=on']::text[], false)
67+
FROM pg_class
68+
WHERE oid = 'public.team_limits'::regclass
69+
`).Scan(&securityInvoker)
70+
require.NoError(t, err)
71+
require.True(t, securityInvoker)
72+
73+
_, err = sqlDB.ExecContext(ctx, `
74+
INSERT INTO public.tiers (
75+
id, name, disk_mb, concurrent_instances, max_length_hours,
76+
default_free_disk_size_mb, max_disk_size_mb
77+
)
78+
VALUES ('en-1038-test', 'EN-1038 test', 10240, 1, 24, 8000, 30000)
79+
`)
80+
require.NoError(t, err)
81+
82+
teamID := uuid.New()
83+
err = db.TestQueries.InsertTestTeam(ctx, testqueries.InsertTestTeamParams{
84+
ID: teamID,
85+
Name: "EN-1038 migration test",
86+
Tier: "en-1038-test",
87+
Email: "en-1038-migration@example.com",
88+
Slug: "en-1038-migration",
89+
})
90+
require.NoError(t, err)
91+
92+
_, err = sqlDB.ExecContext(ctx, `
93+
INSERT INTO public.addons (
94+
team_id, name, extra_disk_mb, extra_max_disk_size_mb, added_by
95+
)
96+
VALUES ($1, 'EN-1038 test add-on', 1000, 3000,
97+
'00000000-0000-0000-0000-000000000000')
98+
`, teamID)
99+
require.NoError(t, err)
100+
101+
var disk, defaultFree, maximum int64
102+
err = sqlDB.QueryRowContext(ctx, `
103+
SELECT disk_mb, default_free_disk_size_mb, max_disk_size_mb
104+
FROM public.team_limits
105+
WHERE id = $1
106+
`, teamID).Scan(&disk, &defaultFree, &maximum)
107+
require.NoError(t, err)
108+
require.Equal(t, int64(11240), disk)
109+
require.Equal(t, int64(9000), defaultFree)
110+
require.Equal(t, int64(33000), maximum)
111+
}

‎packages/db/pkg/testutils/queries/models.go‎

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/db/sqlc.yaml‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,15 @@ overrides:
2121
import: "time"
2222
type: "Time"
2323
pointer: true
24+
# sqlc does not infer nullability for these view expressions.
25+
- column: "public.team_limits.default_free_disk_size_mb"
26+
go_type:
27+
type: "int64"
28+
pointer: true
29+
- column: "public.team_limits.max_disk_size_mb"
30+
go_type:
31+
type: "int64"
32+
pointer: true
2433

2534
sql:
2635
- engine: "postgresql"

0 commit comments

Comments
 (0)