Refactor snapshot to sqlc

packages/api/internal/orchestrator/pause_instance.go

@@ -7,14 +7,15 @@ import (
 	"time"
 
 	"github.com/gogo/status"
+	"github.com/jackc/pgx/v5/pgtype"
 	"google.golang.org/grpc/codes"
 
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/db/types"
-	"github.com/e2b-dev/infra/packages/shared/pkg/db"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 	"github.com/e2b-dev/infra/packages/shared/pkg/models/envbuild"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
@@ -29,39 +30,41 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 	ctx, span := tracer.Start(ctx, "pause-sandbox")
 	defer span.End()
 
-	snapshotConfig := &db.SnapshotInfo{
-		BaseTemplateID:      sbx.BaseTemplateID,
-		SandboxID:           sbx.SandboxID,
-		SandboxStartedAt:    sbx.StartTime,
-		VCPU:                sbx.VCpu,
-		RAMMB:               sbx.RamMB,
-		TotalDiskSizeMB:     sbx.TotalDiskSizeMB,
+	snapshotConfig := queries.UpsertSnapshotParams{
+		// Used if there's no snapshot for this sandbox yet
+		TemplateID:     id.Generate(),
+		TeamID:         sbx.TeamID,
+		BaseTemplateID: sbx.BaseTemplateID,
+		SandboxID:      sbx.SandboxID,
+		StartedAt:      pgtype.Timestamptz{Time: sbx.StartTime, Valid: true},
+		Vcpu:           sbx.VCpu,
+		RamMb:          sbx.RamMB,
+		// We don't know this information
+		FreeDiskSizeMb:      0,
+		TotalDiskSizeMb:     &sbx.TotalDiskSizeMB,
 		Metadata:            sbx.Metadata,
 		KernelVersion:       sbx.KernelVersion,
 		FirecrackerVersion:  sbx.FirecrackerVersion,
-		EnvdVersion:         sbx.EnvdVersion,
-		EnvdSecured:         sbx.EnvdAccessToken != nil,
+		EnvdVersion:         &sbx.EnvdVersion,
+		Secure:              sbx.EnvdAccessToken != nil,
 		AllowInternetAccess: sbx.AllowInternetAccess,
 		AutoPause:           sbx.AutoPause,
 		Config: &types.PausedSandboxConfig{
 			Version: types.PausedSandboxConfigVersion,
 			Network: sbx.Network,
 		},
+		OriginNodeID: node.ID,
+		Status:       string(envbuild.StatusSnapshotting),
 	}
 
-	envBuild, err := o.dbClient.NewSnapshotBuild(
-		ctx,
-		snapshotConfig,
-		sbx.TeamID,
-		sbx.NodeID,
-	)
+	result, err := o.sqlcDB.UpsertSnapshot(ctx, snapshotConfig)
 	if err != nil {
-		telemetry.ReportCriticalError(ctx, "error pausing sandbox", err)
+		telemetry.ReportCriticalError(ctx, "error inserting snapshot for env", err)
 
 		return err
 	}
 
-	err = snapshotInstance(ctx, o, node, sbx, envBuild.EnvID, envBuild.ID.String())
+	err = snapshotInstance(ctx, o, node, sbx, result.TemplateID, result.BuildID.String())
 	if errors.Is(err, PauseQueueExhaustedError{}) {
 		telemetry.ReportCriticalError(ctx, "pause queue exhausted", err)
 
@@ -79,8 +82,8 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 		Status:     string(envbuild.StatusSuccess),
 		FinishedAt: &now,
 		Reason:     types.BuildReason{},
-		BuildID:    envBuild.ID,
-		TemplateID: envBuild.EnvID,
+		BuildID:    result.BuildID,
+		TemplateID: result.TemplateID,
 	})
 	if err != nil {
 		telemetry.ReportCriticalError(ctx, "error pausing sandbox", err)

packages/db/client/tests.go

@@ -2,10 +2,22 @@ package client
 
 import (
 	"context"
+
+	"github.com/jackc/pgx/v5"
 )
 
 func (db *Client) TestsRawSQL(ctx context.Context, sql string, args ...any) error {
 	_, err := db.conn.Exec(ctx, sql, args...)
 
 	return err
 }
+
+func (db *Client) TestsRawSQLQuery(ctx context.Context, sql string, processRows func(pgx.Rows) error, args ...any) error {
+	rows, err := db.conn.Query(ctx, sql, args...)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+
+	return processRows(rows)
+}

packages/db/queries/snapshots/create_new_snapshot.sql

@@ -0,0 +1,75 @@
+-- name: UpsertSnapshot :one
+WITH new_template AS (
+    INSERT INTO "public"."envs" (id, public, created_by, team_id, updated_at)
+    SELECT @template_id, FALSE, NULL, @team_id, now()
+    WHERE NOT EXISTS (
+        SELECT id
+        FROM "public"."snapshots" s
+        WHERE s.sandbox_id = @sandbox_id
+    ) RETURNING id
+),
+
+-- Create a new snapshot or update an existing one
+snapshot as (
+    INSERT INTO "public"."snapshots" (
+       sandbox_id,
+       base_env_id,
+       team_id,
+       env_id,
+       metadata,
+       sandbox_started_at,
+       env_secure,
+       allow_internet_access,
+       origin_node_id,
+       auto_pause,
+       config
+    )
+    VALUES (
+            @sandbox_id,
+            @base_template_id,
+            @team_id,
+            -- If snapshot already exists, new_template id will be null, env_id can't be null, so use placeholder ''
+            COALESCE((SELECT id FROM new_template), ''),
+            @metadata,
+            @started_at,
+            @secure,
+            @allow_internet_access,
+            @origin_node_id,
+            @auto_pause,
+            @config
+   )
+    ON CONFLICT (sandbox_id) DO UPDATE SET
+        metadata = excluded.metadata,
+        sandbox_started_at = excluded.sandbox_started_at,
+        origin_node_id = excluded.origin_node_id,
+        auto_pause = excluded.auto_pause,
+        config = excluded.config
+    RETURNING env_id as template_id
+)
+
+-- Create a new build for the snapshot
+INSERT INTO "public"."env_builds" (
+    env_id,
+    vcpu,
+    ram_mb,
+    free_disk_size_mb,
+    kernel_version,
+    firecracker_version,
+    envd_version,
+    status,
+    cluster_node_id,
+    total_disk_size_mb,
+    updated_at
+) VALUES (
+    (SELECT template_id FROM snapshot),
+    @vcpu,
+    @ram_mb,
+    @free_disk_size_mb,
+    @kernel_version,
+    @firecracker_version,
+    @envd_version,
+    @status,
+    @origin_node_id,
+    @total_disk_size_mb,
+    now()
+) RETURNING id as build_id, env_id as template_id;