Skip to content

Commit 6f513fb

Browse files
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 a84fe31 commit 6f513fb

9 files changed

Lines changed: 248 additions & 5 deletions

File tree

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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 rows that predate the schema expansion in every environment. Keep the column nullable
15+
-- because legacy writers may still omit it during the rollout.
16+
UPDATE "public"."addons"
17+
SET "extra_max_disk_size_mb" = "extra_disk_mb"
18+
WHERE "extra_max_disk_size_mb" IS NULL;
19+
20+
ALTER TABLE "public"."tiers"
21+
ADD CONSTRAINT "tiers_default_free_disk_size_mb_check"
22+
CHECK (default_free_disk_size_mb >= 0),
23+
ADD CONSTRAINT "tiers_max_disk_size_mb_check"
24+
CHECK (max_disk_size_mb > 0),
25+
ADD CONSTRAINT "tiers_default_free_disk_size_lte_max_check"
26+
CHECK (default_free_disk_size_mb <= max_disk_size_mb);
27+
28+
ALTER TABLE "public"."addons"
29+
ADD CONSTRAINT "addons_extra_max_disk_size_mb_check"
30+
CHECK (extra_max_disk_size_mb >= 0);
31+
32+
-- +goose StatementEnd
33+
34+
-- +goose Down
35+
-- +goose StatementBegin
36+
37+
ALTER TABLE "public"."tiers"
38+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_lte_max_check",
39+
DROP CONSTRAINT IF EXISTS "tiers_max_disk_size_mb_check",
40+
DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_mb_check",
41+
DROP COLUMN IF EXISTS "max_disk_size_mb",
42+
DROP COLUMN IF EXISTS "default_free_disk_size_mb";
43+
44+
ALTER TABLE "public"."addons"
45+
DROP CONSTRAINT IF EXISTS "addons_extra_max_disk_size_mb_check",
46+
DROP COLUMN IF EXISTS "extra_max_disk_size_mb";
47+
48+
-- +goose StatementEnd
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.
Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
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 viewColumns string
41+
err = sqlDB.QueryRowContext(ctx, `
42+
SELECT string_agg(column_name, ',' ORDER BY ordinal_position)
43+
FROM information_schema.columns
44+
WHERE table_schema = 'public' AND table_name = 'team_limits'
45+
`).Scan(&viewColumns)
46+
require.NoError(t, err)
47+
48+
const legacyColumns = "id,max_length_hours,concurrent_sandboxes,concurrent_template_builds," +
49+
"max_vcpu,max_ram_mb,disk_mb,events_ttl_days"
50+
require.Equal(t, legacyColumns+",default_free_disk_size_mb,max_disk_size_mb", viewColumns)
51+
52+
var securityInvoker bool
53+
err = sqlDB.QueryRowContext(ctx, `
54+
SELECT COALESCE(reloptions @> ARRAY['security_invoker=on']::text[], false)
55+
FROM pg_class
56+
WHERE oid = 'public.team_limits'::regclass
57+
`).Scan(&securityInvoker)
58+
require.NoError(t, err)
59+
require.True(t, securityInvoker)
60+
61+
_, err = sqlDB.ExecContext(ctx, `
62+
INSERT INTO public.tiers (
63+
id, name, disk_mb, concurrent_instances, max_length_hours,
64+
default_free_disk_size_mb, max_disk_size_mb
65+
)
66+
VALUES ('en-1038-test', 'EN-1038 test', 10240, 1, 24, 8000, 30000)
67+
`)
68+
require.NoError(t, err)
69+
70+
teamID := uuid.New()
71+
err = db.TestQueries.InsertTestTeam(ctx, testqueries.InsertTestTeamParams{
72+
ID: teamID,
73+
Name: "EN-1038 migration test",
74+
Tier: "en-1038-test",
75+
Email: "en-1038-migration@example.com",
76+
Slug: "en-1038-migration",
77+
})
78+
require.NoError(t, err)
79+
80+
_, err = sqlDB.ExecContext(ctx, `
81+
INSERT INTO public.addons (
82+
team_id, name, extra_disk_mb, extra_max_disk_size_mb, added_by
83+
)
84+
VALUES ($1, 'EN-1038 test add-on', 1000, 3000,
85+
'00000000-0000-0000-0000-000000000000')
86+
`, teamID)
87+
require.NoError(t, err)
88+
89+
var disk, defaultFree, maximum int64
90+
err = sqlDB.QueryRowContext(ctx, `
91+
SELECT disk_mb, default_free_disk_size_mb, max_disk_size_mb
92+
FROM public.team_limits
93+
WHERE id = $1
94+
`, teamID).Scan(&disk, &defaultFree, &maximum)
95+
require.NoError(t, err)
96+
require.Equal(t, int64(11240), disk)
97+
require.Equal(t, int64(9000), defaultFree)
98+
require.Equal(t, int64(33000), maximum)
99+
}

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)