Simplify the query

Author: jakubno

packages/api/internal/orchestrator/pause_instance.go

@@ -30,7 +30,7 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 	ctx, span := tracer.Start(ctx, "pause-sandbox")
 	defer span.End()
 
-	snapshotConfig := queries.UpsertSnapshotEnvAndBuildParams{
+	snapshotConfig := queries.UpsertSnapshotParams{
 		// Used if there's no snapshot for this sandbox yet
 		TemplateID:          id.Generate(),
 		TeamID:              sbx.TeamID,
@@ -55,7 +55,7 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 		Status:       string(envbuild.StatusSnapshotting),
 	}
 
-	result, err := o.sqlcDB.UpsertSnapshotEnvAndBuild(ctx, snapshotConfig)
+	result, err := o.sqlcDB.UpsertSnapshot(ctx, snapshotConfig)
 	if err != nil {
 		telemetry.ReportCriticalError(ctx, "error inserting snapshot for env", err)
 

packages/db/queries/create_new_snapshot.sql.go

@@ -1,30 +1,17 @@
--- name: UpsertSnapshotEnvAndBuild :one
--- Try to update an existing snapshot
-WITH updated_snapshot AS (
-    UPDATE "public"."snapshots" s
-    SET
-        metadata           = @metadata,
-        sandbox_started_at = @started_at,
-        origin_node_id     = @origin_node_id,
-        auto_pause         = @auto_pause,
-        config             = @config
-    FROM "public"."envs" e
-    WHERE
-        s.sandbox_id = @sandbox_id
-            AND e.id      = s.env_id
-            AND e.team_id = @team_id
-    RETURNING s.env_id
-),
--- otherwise insert a new one with a new env
- inserted_env 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 1 FROM updated_snapshot)
-     ON CONFLICT (id) DO NOTHING
-     RETURNING id AS env_id
+-- 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
 ),
- inserted_snapshot AS (
-     INSERT INTO "public"."snapshots" (
+
+-- Create a new snapshot or update an existing one
+snapshot as (
+    INSERT INTO "public"."snapshots" (
        sandbox_id,
        base_env_id,
        team_id,
@@ -36,32 +23,28 @@ WITH updated_snapshot AS (
        origin_node_id,
        auto_pause,
        config
-     )
-     SELECT
-         @sandbox_id,
-         @base_template_id,
-         @team_id,
-         COALESCE(
-                 (SELECT env_id FROM inserted_env LIMIT 1),
-                 @template_id      -- env already existed
-         ) AS env_id,
-         @metadata,
-         @started_at,
-         @secure,
-         @allow_internet_access,
-         @origin_node_id,
-         @auto_pause,
-         @config
-     WHERE NOT EXISTS (SELECT 1 FROM updated_snapshot)
-     RETURNING env_id
- ),
- -- If we updated an existing snapshot, use that env_id.
- -- Otherwise use env_id from the newly inserted snapshot.
- final_env AS (
-     SELECT env_id FROM updated_snapshot
-     UNION ALL
-     SELECT env_id FROM inserted_snapshot
- )
+    )
+    VALUES (
+            @sandbox_id,
+            @base_template_id,
+            @team_id,
+            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 = @metadata,
+        sandbox_started_at = @started_at,
+        origin_node_id = @origin_node_id,
+        auto_pause = @auto_pause,
+        config = @config
+    RETURNING env_id as template_id
+)
 
 -- Create a new build for the snapshot
 INSERT INTO "public"."env_builds" (
@@ -76,9 +59,8 @@ INSERT INTO "public"."env_builds" (
     cluster_node_id,
     total_disk_size_mb,
     updated_at
-)
-SELECT
-    (SELECT env_id FROM final_env LIMIT 1),
+) VALUES (
+    (SELECT template_id FROM snapshot),
     @vcpu,
     @ram_mb,
     0,
@@ -89,4 +71,4 @@ SELECT
     @origin_node_id,
     @total_disk_size_mb,
     now()
-RETURNING id as build_id, env_id as template_id;
+) RETURNING id as build_id, env_id as template_id;

packages/db/queries/snapshots/create_new_snapshot.sql

@@ -33,6 +33,7 @@ func NewClient() (*Client, error) {
 			"",
 			ldclient.Config{
 				DataSource: LaunchDarklyOfflineStore,
+				Offline:    true,
 			},
 			0)
 		if err != nil {