[EN-1038] feat(db): add template disk entitlement schema

packages/db/migrations/20260714091414_add_template_disk_entitlements.sql

@@ -0,0 +1,53 @@
+-- +goose Up
+-- +goose StatementBegin
+
+ALTER TABLE "public"."tiers"
+    -- Free-rootfs target in the existing MiB convention.
+    ADD COLUMN IF NOT EXISTS "default_free_disk_size_mb" bigint,
+    -- Total logical-rootfs ceiling before active add-ons.
+    ADD COLUMN IF NOT EXISTS "max_disk_size_mb" bigint;
+
+ALTER TABLE "public"."addons"
+    -- Total-ceiling increment in the existing MiB convention.
+    ADD COLUMN IF NOT EXISTS "extra_max_disk_size_mb" bigint;
+
+-- Populate existing tiers in every environment before the effective limits are exposed.
+UPDATE "public"."tiers"
+SET "default_free_disk_size_mb" = "disk_mb",
+    "max_disk_size_mb" = "disk_mb" + 25000;
+
+-- Populate rows that predate the schema expansion in every environment. Keep the column nullable
+-- because legacy writers may still omit it during the rollout.
+UPDATE "public"."addons"
+SET "extra_max_disk_size_mb" = "extra_disk_mb"
+WHERE "extra_max_disk_size_mb" IS NULL;
+
+ALTER TABLE "public"."tiers"
+    ADD CONSTRAINT "tiers_default_free_disk_size_mb_check"
+        CHECK (default_free_disk_size_mb >= 0),
+    ADD CONSTRAINT "tiers_max_disk_size_mb_check"
+        CHECK (max_disk_size_mb > 0),
+    ADD CONSTRAINT "tiers_default_free_disk_size_lte_max_check"
+        CHECK (default_free_disk_size_mb <= max_disk_size_mb);
+
+ALTER TABLE "public"."addons"
+    ADD CONSTRAINT "addons_extra_max_disk_size_mb_check"
+        CHECK (extra_max_disk_size_mb >= 0);
+
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+
+ALTER TABLE "public"."tiers"
+    DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_lte_max_check",
+    DROP CONSTRAINT IF EXISTS "tiers_max_disk_size_mb_check",
+    DROP CONSTRAINT IF EXISTS "tiers_default_free_disk_size_mb_check",
+    DROP COLUMN IF EXISTS "max_disk_size_mb",
+    DROP COLUMN IF EXISTS "default_free_disk_size_mb";
+
+ALTER TABLE "public"."addons"
+    DROP CONSTRAINT IF EXISTS "addons_extra_max_disk_size_mb_check",
+    DROP COLUMN IF EXISTS "extra_max_disk_size_mb";
+
+-- +goose StatementEnd

packages/db/migrations/20260714091415_extend_team_limits.sql

@@ -0,0 +1,68 @@
+-- +goose Up
+-- +goose StatementBegin
+
+CREATE OR REPLACE VIEW "public"."team_limits"
+WITH (security_invoker=on) AS
+SELECT
+    t.id,
+    tier.max_length_hours,
+    (tier.concurrent_instances + a.extra_concurrent_sandboxes) AS concurrent_sandboxes,
+    (tier.concurrent_template_builds + a.extra_concurrent_template_builds) AS concurrent_template_builds,
+    (tier.max_vcpu + a.extra_max_vcpu) AS max_vcpu,
+    (tier.max_ram_mb + a.extra_max_ram_mb) AS max_ram_mb,
+    (tier.disk_mb + a.extra_disk_mb) AS disk_mb,
+    (tier.events_ttl_days + a.extra_events_ttl_days) AS events_ttl_days,
+    (tier.default_free_disk_size_mb + a.extra_disk_mb)::bigint AS default_free_disk_size_mb,
+    (tier.max_disk_size_mb + a.extra_max_disk_size_mb)::bigint AS max_disk_size_mb
+FROM "public"."teams" t
+JOIN "public"."tiers" tier ON t.tier = tier.id
+LEFT JOIN LATERAL (
+    SELECT COALESCE(SUM(extra_concurrent_sandboxes), 0)::bigint AS extra_concurrent_sandboxes,
+           COALESCE(SUM(extra_concurrent_template_builds), 0)::bigint AS extra_concurrent_template_builds,
+           COALESCE(SUM(extra_max_vcpu), 0)::bigint AS extra_max_vcpu,
+           COALESCE(SUM(extra_max_ram_mb), 0)::bigint AS extra_max_ram_mb,
+           COALESCE(SUM(extra_disk_mb), 0)::bigint AS extra_disk_mb,
+           COALESCE(SUM(extra_events_ttl_days), 0)::bigint AS extra_events_ttl_days,
+           COALESCE(SUM(COALESCE(extra_max_disk_size_mb, extra_disk_mb)), 0)::bigint AS extra_max_disk_size_mb
+    FROM "public"."addons" addon
+    WHERE addon.team_id = t.id
+      AND addon.valid_from <= now()
+      AND (addon.valid_to IS NULL OR addon.valid_to > now())
+) a ON true;
+
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+
+-- PostgreSQL cannot remove appended view columns with CREATE OR REPLACE, so restore the previous
+-- definition transactionally. Concurrent readers cannot observe the view between these statements.
+DROP VIEW IF EXISTS "public"."team_limits";
+
+CREATE VIEW "public"."team_limits"
+WITH (security_invoker=on) AS
+SELECT
+    t.id,
+    tier.max_length_hours,
+    (tier.concurrent_instances + a.extra_concurrent_sandboxes) AS concurrent_sandboxes,
+    (tier.concurrent_template_builds + a.extra_concurrent_template_builds) AS concurrent_template_builds,
+    (tier.max_vcpu + a.extra_max_vcpu) AS max_vcpu,
+    (tier.max_ram_mb + a.extra_max_ram_mb) AS max_ram_mb,
+    (tier.disk_mb + a.extra_disk_mb) AS disk_mb,
+    (tier.events_ttl_days + a.extra_events_ttl_days) AS events_ttl_days
+FROM "public"."teams" t
+JOIN "public"."tiers" tier ON t.tier = tier.id
+LEFT JOIN LATERAL (
+    SELECT COALESCE(SUM(extra_concurrent_sandboxes), 0)::bigint AS extra_concurrent_sandboxes,
+           COALESCE(SUM(extra_concurrent_template_builds), 0)::bigint AS extra_concurrent_template_builds,
+           COALESCE(SUM(extra_max_vcpu), 0)::bigint AS extra_max_vcpu,
+           COALESCE(SUM(extra_max_ram_mb), 0)::bigint AS extra_max_ram_mb,
+           COALESCE(SUM(extra_disk_mb), 0)::bigint AS extra_disk_mb,
+           COALESCE(SUM(extra_events_ttl_days), 0)::bigint AS extra_events_ttl_days
+    FROM "public"."addons" addon
+    WHERE addon.team_id = t.id
+      AND addon.valid_from <= now()
+      AND (addon.valid_to IS NULL OR addon.valid_to > now())
+) a ON true;
+
+-- +goose StatementEnd

packages/db/pkg/tests/disk_entitlements_migration_test.go

@@ -0,0 +1,111 @@
+package tests
+
+import (
+	"database/sql"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	testqueries "github.com/e2b-dev/infra/packages/db/pkg/testutils/queries"
+)
+
+func TestDiskEntitlementsMigration(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+
+	sqlDB, err := sql.Open("pgx", db.ConnStr())
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
+
+	var columnCount int64
+	var nullableWithoutDefaults bool
+	err = sqlDB.QueryRowContext(ctx, `
+		SELECT COUNT(*), BOOL_AND(data_type = 'bigint' AND is_nullable = 'YES' AND column_default IS NULL)
+		FROM information_schema.columns
+		WHERE table_schema = 'public'
+		  AND (table_name, column_name) IN (
+			('tiers', 'default_free_disk_size_mb'),
+			('tiers', 'max_disk_size_mb'),
+			('addons', 'extra_max_disk_size_mb')
+		  )
+	`).Scan(&columnCount, &nullableWithoutDefaults)
+	require.NoError(t, err)
+	require.Equal(t, int64(3), columnCount)
+	require.True(t, nullableWithoutDefaults)
+
+	var tierCount, invalidTierCount int64
+	err = sqlDB.QueryRowContext(ctx, `
+		SELECT COUNT(*), COUNT(*) FILTER (
+			WHERE default_free_disk_size_mb IS DISTINCT FROM disk_mb
+			   OR max_disk_size_mb IS DISTINCT FROM disk_mb + 25000
+		)
+		FROM public.tiers
+	`).Scan(&tierCount, &invalidTierCount)
+	require.NoError(t, err)
+	require.NotZero(t, tierCount)
+	require.Zero(t, invalidTierCount)
+
+	var viewColumns string
+	err = sqlDB.QueryRowContext(ctx, `
+		SELECT string_agg(column_name, ',' ORDER BY ordinal_position)
+		FROM information_schema.columns
+		WHERE table_schema = 'public' AND table_name = 'team_limits'
+	`).Scan(&viewColumns)
+	require.NoError(t, err)
+
+	const legacyColumns = "id,max_length_hours,concurrent_sandboxes,concurrent_template_builds," +
+		"max_vcpu,max_ram_mb,disk_mb,events_ttl_days"
+	require.Equal(t, legacyColumns+",default_free_disk_size_mb,max_disk_size_mb", viewColumns)
+
+	var securityInvoker bool
+	err = sqlDB.QueryRowContext(ctx, `
+		SELECT COALESCE(reloptions @> ARRAY['security_invoker=on']::text[], false)
+		FROM pg_class
+		WHERE oid = 'public.team_limits'::regclass
+	`).Scan(&securityInvoker)
+	require.NoError(t, err)
+	require.True(t, securityInvoker)
+
+	_, err = sqlDB.ExecContext(ctx, `
+		INSERT INTO public.tiers (
+			id, name, disk_mb, concurrent_instances, max_length_hours,
+			default_free_disk_size_mb, max_disk_size_mb
+		)
+		VALUES ('en-1038-test', 'EN-1038 test', 10240, 1, 24, 8000, 30000)
+	`)
+	require.NoError(t, err)
+
+	teamID := uuid.New()
+	err = db.TestQueries.InsertTestTeam(ctx, testqueries.InsertTestTeamParams{
+		ID:    teamID,
+		Name:  "EN-1038 migration test",
+		Tier:  "en-1038-test",
+		Email: "en-1038-migration@example.com",
+		Slug:  "en-1038-migration",
+	})
+	require.NoError(t, err)
+
+	_, err = sqlDB.ExecContext(ctx, `
+		INSERT INTO public.addons (
+			team_id, name, extra_disk_mb, extra_max_disk_size_mb, added_by
+		)
+		VALUES ($1, 'EN-1038 test add-on', 1000, 3000,
+			'00000000-0000-0000-0000-000000000000')
+	`, teamID)
+	require.NoError(t, err)
+
+	var disk, defaultFree, maximum int64
+	err = sqlDB.QueryRowContext(ctx, `
+		SELECT disk_mb, default_free_disk_size_mb, max_disk_size_mb
+		FROM public.team_limits
+		WHERE id = $1
+	`, teamID).Scan(&disk, &defaultFree, &maximum)
+	require.NoError(t, err)
+	require.Equal(t, int64(11240), disk)
+	require.Equal(t, int64(9000), defaultFree)
+	require.Equal(t, int64(33000), maximum)
+}

packages/db/sqlc.yaml

@@ -21,6 +21,15 @@ overrides:
           import: "time"
           type: "Time"
           pointer: true
+      # sqlc does not infer nullability for these view expressions.
+      - column: "public.team_limits.default_free_disk_size_mb"
+        go_type:
+          type: "int64"
+          pointer: true
+      - column: "public.team_limits.max_disk_size_mb"
+        go_type:
+          type: "int64"
+          pointer: true
 
 sql:
   - engine: "postgresql"