feat: allow m n builds with tags

Author: dobrac

packages/api/internal/cache/templates/cache.go

@@ -2,10 +2,10 @@ package templatecache
 
 import (
 	"context"
-	"database/sql"
 	"errors"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -14,8 +14,10 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/api"
 	"github.com/e2b-dev/infra/packages/api/internal/utils"
 	sqlcdb "github.com/e2b-dev/infra/packages/db/client"
+	"github.com/e2b-dev/infra/packages/db/dberrors"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/cache"
+	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 )
 
 const (
@@ -29,6 +31,7 @@ type TemplateInfo struct {
 	teamID    uuid.UUID
 	clusterID uuid.UUID
 	build     *queries.EnvBuild
+	tag       *string
 }
 
 type AliasCache struct {
@@ -72,7 +75,7 @@ func NewTemplateCache(db *sqlcdb.Client) *TemplateCache {
 		RefreshTimeout:  refreshTimeout,
 		// With this we can use alias for getting template info without having it as a key in the cache
 		ExtractKeyFunc: func(value *TemplateInfo) string {
-			return value.template.TemplateID
+			return buildCacheKey(value.template.TemplateID, value.tag)
 		},
 	}
 	aliasCache := NewAliasCache()
@@ -84,15 +87,25 @@ func NewTemplateCache(db *sqlcdb.Client) *TemplateCache {
 	}
 }
 
-func (c *TemplateCache) Get(ctx context.Context, aliasOrEnvID string, teamID uuid.UUID, clusterID uuid.UUID, public bool) (*api.Template, *queries.EnvBuild, *api.APIError) {
+func buildCacheKey(templateID string, tag *string) string {
+	if tag == nil {
+		return templateID + ":" + id.DefaultTag
+	}
+
+	return templateID + ":" + *tag
+}
+
+func (c *TemplateCache) Get(ctx context.Context, aliasOrEnvID string, tag *string, teamID uuid.UUID, clusterID uuid.UUID, public bool) (*api.Template, *queries.EnvBuild, *api.APIError) {
 	// Resolve alias to template ID if needed
 	templateID, found := c.aliasCache.Get(aliasOrEnvID)
 	if !found {
 		templateID = aliasOrEnvID
 	}
 
+	cacheKey := buildCacheKey(templateID, tag)
+
 	// Fetch or get from cache with automatic refresh
-	templateInfo, err := c.cache.GetOrSet(ctx, templateID, c.fetchTemplateInfo)
+	templateInfo, err := c.cache.GetOrSet(ctx, cacheKey, c.fetchTemplateInfo)
 	if err != nil {
 		var apiErr *api.APIError
 		if errors.As(err, &apiErr) {
@@ -115,11 +128,24 @@ func (c *TemplateCache) Get(ctx context.Context, aliasOrEnvID string, teamID uui
 }
 
 // fetchTemplateInfo fetches template info from the database
-func (c *TemplateCache) fetchTemplateInfo(ctx context.Context, aliasOrEnvID string) (*TemplateInfo, error) {
-	result, err := c.db.GetTemplateWithBuild(ctx, aliasOrEnvID)
+func (c *TemplateCache) fetchTemplateInfo(ctx context.Context, cacheKey string) (*TemplateInfo, error) {
+	aliasOrEnvID, tag, err := id.ParseTemplateIDOrAliasWithTag(cacheKey)
 	if err != nil {
-		if errors.Is(err, sql.ErrNoRows) {
-			return nil, &api.APIError{Code: http.StatusNotFound, ClientMsg: fmt.Sprintf("template '%s' not found", aliasOrEnvID), Err: err}
+		return nil, &api.APIError{Code: http.StatusBadRequest, ClientMsg: fmt.Sprintf("invalid template ID: %s", err), Err: err}
+	}
+
+	result, err := c.db.GetTemplateWithBuildByTag(ctx, queries.GetTemplateWithBuildByTagParams{
+		AliasOrEnvID: aliasOrEnvID,
+		Tag:          tag,
+	})
+	if err != nil {
+		if dberrors.IsNotFoundError(err) {
+			tagMsg := ""
+			if tag != nil {
+				tagMsg = fmt.Sprintf(" with tag '%s'", *tag)
+			}
+
+			return nil, &api.APIError{Code: http.StatusNotFound, ClientMsg: fmt.Sprintf("template '%s'%s not found", aliasOrEnvID, tagMsg), Err: err}
 		}
 
 		return nil, &api.APIError{Code: http.StatusInternalServerError, ClientMsg: fmt.Sprintf("error while getting template: %v", err), Err: err}
@@ -129,7 +155,7 @@ func (c *TemplateCache) fetchTemplateInfo(ctx context.Context, aliasOrEnvID stri
 	template := result.Env
 	clusterID := utils.WithClusterFallback(template.ClusterID)
 
-	// Update alias cache
+	// Update alias cache (without tag, as aliases map to template IDs)
 	c.aliasCache.Set(template.ID, template.ID)
 	for _, alias := range result.Aliases {
 		c.aliasCache.Set(alias, template.ID)
@@ -145,12 +171,22 @@ func (c *TemplateCache) fetchTemplateInfo(ctx context.Context, aliasOrEnvID stri
 		teamID:    template.TeamID,
 		clusterID: clusterID,
 		build:     build,
+		tag:       tag,
 	}, nil
 }
 
-// Invalidate invalidates the cache for the given templateID
-func (c *TemplateCache) Invalidate(templateID string) {
-	c.cache.Delete(templateID)
+func (c *TemplateCache) Invalidate(templateID string, tag *string) {
+	c.cache.Delete(buildCacheKey(templateID, tag))
+}
+
+// Invalidate invalidates the cache for the given templateID across all tags
+func (c *TemplateCache) InvalidateAllTags(templateID string) {
+	templateIDKey := templateID + ":"
+	for _, key := range c.cache.Keys() {
+		if strings.HasPrefix(key, templateIDKey) {
+			c.cache.Delete(key)
+		}
+	}
 }
 
 func (c *TemplateCache) Close(ctx context.Context) error {

packages/api/internal/handlers/deprecated_template_request_build.go

@@ -84,7 +84,7 @@ func (a *APIStore) PostTemplatesTemplateID(c *gin.Context, rawTemplateID api.Tem
 		return
 	}
 
-	templateID, err := id.CleanTemplateID(rawTemplateID)
+	templateID, _, err := id.ParseTemplateIDOrAliasWithTag(rawTemplateID)
 	if err != nil {
 		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Invalid template ID: %s", rawTemplateID))
 		telemetry.ReportCriticalError(c.Request.Context(), "invalid template ID", err)

packages/api/internal/handlers/sandbox_create.go

@@ -66,23 +66,24 @@ func (a *APIStore) PostSandboxes(c *gin.Context) {
 
 	telemetry.ReportEvent(ctx, "Parsed body")
 
-	cleanedAliasOrEnvID, err := id.CleanTemplateID(body.TemplateID)
+	// Parse template ID and optional tag in the format "templateID:tag"
+	cleanedAliasOrEnvID, tag, err := id.ParseTemplateIDOrAliasWithTag(body.TemplateID)
 	if err != nil {
-		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Invalid environment ID: %s", err))
+		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Invalid template ID: %s", err))
 
-		telemetry.ReportCriticalError(ctx, "error when cleaning env ID", err)
+		telemetry.ReportCriticalError(ctx, "error when parsing template ID", err)
 
 		return
 	}
 
-	telemetry.ReportEvent(ctx, "Cleaned template ID")
+	telemetry.ReportEvent(ctx, "Parsed template ID and tag")
 
 	_, templateSpan := tracer.Start(ctx, "get-template")
 	defer templateSpan.End()
 
 	// Check if team has access to the environment
 	clusterID := utils.WithClusterFallback(teamInfo.Team.ClusterID)
-	env, build, checkErr := a.templateCache.Get(ctx, cleanedAliasOrEnvID, teamInfo.Team.ID, clusterID, true)
+	env, build, checkErr := a.templateCache.Get(ctx, cleanedAliasOrEnvID, tag, teamInfo.Team.ID, clusterID, true)
 	if checkErr != nil {
 		telemetry.ReportCriticalError(ctx, "error when getting template", checkErr.Err)
 		a.sendAPIStoreError(c, checkErr.Code, checkErr.ClientMsg)

packages/api/internal/handlers/sandbox_kill.go

@@ -66,7 +66,7 @@ func (a *APIStore) deleteSnapshot(ctx context.Context, sandboxID string, teamID
 		}
 	}(context.WithoutCancel(ctx))
 
-	a.templateCache.Invalidate(snapshot.TemplateID)
+	a.templateCache.Invalidate(snapshot.TemplateID, nil)
 
 	return nil
 }

packages/api/internal/handlers/template_delete.go

@@ -20,7 +20,7 @@ import (
 func (a *APIStore) DeleteTemplatesTemplateID(c *gin.Context, aliasOrTemplateID api.TemplateID) {
 	ctx := c.Request.Context()
 
-	cleanedAliasOrTemplateID, err := id.CleanTemplateID(aliasOrTemplateID)
+	cleanedAliasOrTemplateID, _, err := id.ParseTemplateIDOrAliasWithTag(aliasOrTemplateID)
 	if err != nil {
 		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Invalid template ID: %s", aliasOrTemplateID))
 
@@ -119,7 +119,7 @@ func (a *APIStore) DeleteTemplatesTemplateID(c *gin.Context, aliasOrTemplateID a
 		telemetry.ReportEvent(ctx, "deleted template from storage")
 	}
 
-	a.templateCache.Invalidate(templateID)
+	a.templateCache.InvalidateAllTags(templateID)
 
 	telemetry.ReportEvent(ctx, "deleted template from db")
 

packages/api/internal/handlers/template_update.go

@@ -28,7 +28,7 @@ func (a *APIStore) PatchTemplatesTemplateID(c *gin.Context, aliasOrTemplateID ap
 		return
 	}
 
-	cleanedAliasOrTemplateID, err := id.CleanTemplateID(aliasOrTemplateID)
+	cleanedAliasOrTemplateID, tag, err := id.ParseTemplateIDOrAliasWithTag(aliasOrTemplateID)
 	if err != nil {
 		a.sendAPIStoreError(c, http.StatusBadRequest, fmt.Sprintf("Invalid template ID: %s", aliasOrTemplateID))
 
@@ -95,7 +95,7 @@ func (a *APIStore) PatchTemplatesTemplateID(c *gin.Context, aliasOrTemplateID ap
 		}
 	}
 
-	a.templateCache.Invalidate(template.ID)
+	a.templateCache.Invalidate(template.ID, tag)
 
 	telemetry.ReportEvent(ctx, "updated template")
 

packages/api/internal/template-manager/create_template.go

@@ -189,7 +189,7 @@ func (tm *TemplateManager) CreateTemplate(
 		}
 
 		// Invalidate the cache
-		tm.templateCache.Invalidate(templateID)
+		tm.templateCache.Invalidate(templateID, nil)
 	}(context.WithoutCancel(ctx))
 
 	return nil

packages/api/internal/template/register_build.go

@@ -128,7 +128,7 @@ func RegisterBuild(
 
 	var alias string
 	if data.Alias != nil {
-		alias, err = id.CleanTemplateID(*data.Alias)
+		alias, _, err = id.ParseTemplateIDOrAliasWithTag(*data.Alias)
 		if err != nil {
 			telemetry.ReportCriticalError(ctx, "invalid alias", err)
 

packages/db/migrations/20251126120000_allow_m_n_builds_with_tags.sql

@@ -0,0 +1,91 @@
+-- +goose Up
+-- +goose StatementBegin
+-- 1. Create the new env_build_assignments table
+CREATE TABLE env_build_assignments (
+    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
+    env_id TEXT NOT NULL,
+    build_id UUID NOT NULL,
+    tag TEXT,
+    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
+    
+    -- Add foreign key constraints
+    CONSTRAINT fk_env_build_assignments_env 
+        FOREIGN KEY (env_id) 
+        REFERENCES envs(id) 
+        ON DELETE CASCADE,
+    
+    CONSTRAINT fk_env_build_assignments_build 
+        FOREIGN KEY (build_id) 
+        REFERENCES env_builds(id) 
+        ON DELETE CASCADE,
+    
+    -- Add unique constraint to prevent duplicate assignments for the same build+tag combination
+    -- This ensures data consistency during migration and normal operations
+    CONSTRAINT uq_env_build_assignments_build_tag 
+        UNIQUE (build_id, tag, created_at)
+);
+ALTER TABLE "public"."env_build_assignments" ENABLE ROW LEVEL SECURITY;
+-- +goose StatementEnd
+
+-- +goose StatementBegin
+-- 2. Create trigger function to sync env_builds.env_id changes to env_build_assignments
+-- This table is append-only for 'latest' tag. Queries use ORDER BY created_at DESC LIMIT 1 to get current assignment.
+CREATE OR REPLACE FUNCTION sync_env_build_assignment()
+RETURNS TRIGGER AS $$
+BEGIN
+    -- On INSERT or UPDATE, if env_id is set, append new assignment with 'latest' tag
+    IF NEW.env_id IS NOT NULL THEN
+        -- Check if this is actually a change (for UPDATE operations)
+        IF TG_OP = 'INSERT' OR (TG_OP = 'UPDATE' AND (OLD.env_id IS NULL OR OLD.env_id != NEW.env_id)) THEN
+            -- Append new assignment with 'latest' tag (append-only)
+            -- Use ON CONFLICT DO NOTHING to handle race conditions during migration
+            INSERT INTO env_build_assignments (env_id, build_id, tag, created_at)
+            VALUES (NEW.env_id, NEW.id, 'latest', CURRENT_TIMESTAMP)
+            ON CONFLICT (build_id, tag, created_at) DO NOTHING;
+        END IF;
+    END IF;
+    
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+-- +goose StatementEnd
+
+-- +goose StatementBegin
+-- 3. Create trigger to automatically sync changes (BEFORE data migration to prevent race conditions)
+CREATE TRIGGER trigger_sync_env_build_assignment
+    AFTER INSERT OR UPDATE ON env_builds
+    FOR EACH ROW
+    EXECUTE FUNCTION sync_env_build_assignment();
+-- +goose StatementEnd
+
+-- +goose StatementBegin
+-- 4. Migrate existing data from direct relationship
+-- The env_builds table has env_id column pointing to envs
+-- The trigger is already active, so ON CONFLICT handles any concurrent inserts
+INSERT INTO env_build_assignments (env_id, build_id, tag, created_at)
+SELECT 
+    env_id,
+    id as build_id,
+    'latest' as tag,
+    created_at
+FROM env_builds
+WHERE env_id IS NOT NULL
+ON CONFLICT (build_id, tag, created_at) DO NOTHING;
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+-- 1. Drop the trigger
+DROP TRIGGER IF EXISTS trigger_sync_env_build_assignment ON env_builds;
+-- +goose StatementEnd
+
+-- +goose StatementBegin
+-- 2. Drop the trigger function
+DROP FUNCTION IF EXISTS sync_env_build_assignment();
+-- +goose StatementEnd
+
+-- +goose StatementBegin
+-- 3. Drop the assignments table
+DROP TABLE IF EXISTS env_build_assignments;
+-- +goose StatementEnd
+