fix(db): set a build's env and team when it is inserted

GitOrigin-RevId: 6adcdd1e7add73fa382d0a3cc7d65ab047873740

Author: charlie-e2b

packages/api/internal/template/register_build.go

@@ -287,6 +287,7 @@ func registerBuild(
 	// used in TemplateBuildMetadata and SetFinished overwrites these rows.
 	err = client.CreateTemplateBuild(ctx, queries.CreateTemplateBuildParams{
 		BuildID:            buildID,
+		TemplateID:         data.TemplateID,
 		Status:             dbtypes.BuildStatusWaiting,
 		RamMb:              ramMB,
 		Vcpu:               cpuCount,

packages/db/pkg/tests/build_owner_test.go

@@ -0,0 +1,65 @@
+package tests
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/require"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+	"github.com/e2b-dev/infra/packages/db/pkg/types"
+	"github.com/e2b-dev/infra/packages/db/queries"
+)
+
+// Template builds and snapshot builds set env_builds.env_id and team_id when
+// they are inserted, so the build is complete without the backfill triggers
+// on env_build_assignments. The triggers are disabled here to prove it.
+func TestBuildInsertsSetEnvAndTeamWithoutTriggers(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	ctx := t.Context()
+	require.NoError(t, db.SqlcClient.TestsRawSQL(ctx, `ALTER TABLE public.env_build_assignments DISABLE TRIGGER USER`))
+
+	teamID := testutils.CreateTestTeam(t, db)
+	templateID := testutils.CreateTestTemplate(t, db, teamID)
+	read := func(t *testing.T, buildID uuid.UUID) (*string, *uuid.UUID) {
+		t.Helper()
+		var envID *string
+		var buildTeamID *uuid.UUID
+		require.NoError(t, db.SqlcClient.TestsRawSQLQuery(ctx, `SELECT env_id, team_id FROM public.env_builds WHERE id = $1`,
+			func(rows pgx.Rows) error {
+				require.True(t, rows.Next())
+
+				return rows.Scan(&envID, &buildTeamID)
+			}, buildID))
+
+		return envID, buildTeamID
+	}
+
+	buildID := uuid.New()
+	require.NoError(t, db.SqlcClient.CreateTemplateBuild(ctx, queries.CreateTemplateBuildParams{
+		BuildID:            buildID,
+		TemplateID:         templateID,
+		Status:             types.BuildStatusWaiting,
+		RamMb:              512,
+		Vcpu:               1,
+		KernelVersion:      "6.1.0",
+		FirecrackerVersion: "1.4.0",
+		FreeDiskSizeMb:     512,
+	}))
+	envID, buildTeamID := read(t, buildID)
+	require.Equal(t, &templateID, envID)
+	require.Equal(t, &teamID, buildTeamID)
+
+	// The first pause creates the snapshot env; a second pause of the same
+	// sandbox adds a build to the existing one.
+	sandboxID := "sbx-" + uuid.NewString()
+	for range 2 {
+		snapshot := testutils.UpsertTestSnapshot(t, ctx, db, "snap-"+uuid.NewString(), sandboxID, teamID, templateID)
+		envID, buildTeamID := read(t, snapshot.BuildID)
+		require.Equal(t, &snapshot.TemplateID, envID)
+		require.Equal(t, &teamID, buildTeamID)
+	}
+}

packages/db/queries/create_new_snapshot.sql.go

@@ -56,6 +56,8 @@ snapshot as (
 -- the row is always inserted (CPU info is NULL if the source build is missing).
 new_build as (
     INSERT INTO "public"."env_builds" (
+        env_id,
+        team_id,
         vcpu,
         ram_mb,
         free_disk_size_mb,
@@ -73,6 +75,8 @@ new_build as (
         cpu_flags
     )
     VALUES (
+        (SELECT template_id FROM snapshot),
+        @team_id,
         @vcpu,
         @ram_mb,
         @free_disk_size_mb,

packages/db/queries/create_template.sql.go

@@ -32,8 +32,12 @@ WHERE build_id IN (SELECT id FROM invalidated);
 -- completes. The template-manager reports the versions it actually used via
 -- TemplateBuildMetadata, and FinishTemplateBuild overwrites these fields with
 -- the reported values.
+-- env_id and team_id are set here rather than left to the trigger on
+-- env_build_assignments, so the build is complete without it.
 INSERT INTO "public"."env_builds" (
     id,
+    env_id,
+    team_id,
     updated_at,
     status,
     ram_mb,
@@ -47,6 +51,8 @@ INSERT INTO "public"."env_builds" (
     version
 ) VALUES (
     @build_id,
+    sqlc.arg(template_id)::text,
+    (SELECT team_id FROM "public"."envs" WHERE id = sqlc.arg(template_id)::text),
     NOW(),
     @status,
     @ram_mb,