feat: add template versioning using tags (#1524)

Author: dobrac

packages/api/internal/api/api.gen.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/api/spec.gen.go

@@ -14,7 +14,6 @@ import (
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/db/testutils"
 	"github.com/e2b-dev/infra/packages/db/types"
-	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 // createTestTeam creates a test team in the database using raw SQL
@@ -85,7 +84,7 @@ func createTestSnapshot(t *testing.T, db *client.Client, teamID uuid.UUID, baseE
 				},
 			},
 		},
-		OriginNodeID: utils.ToPtr("node-1"),
+		OriginNodeID: "node-1",
 		Status:       "success",
 	}
 

packages/api/internal/api/types.gen.go

@@ -85,7 +85,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)
@@ -156,14 +156,44 @@ func (a *APIStore) buildTemplate(
 ) (*template.RegisterBuildResponse, *api.APIError) {
 	firecrackerVersion := a.featureFlags.StringFlag(ctx, featureflags.BuildFirecrackerVersion)
 
+	var alias *string
+	var tags []string
+
+	if body.Alias != nil {
+		var err error
+		a, t, err := id.ParseTemplateIDOrAliasWithTag(*body.Alias)
+		if err != nil {
+			return nil, &api.APIError{
+				Code:      http.StatusBadRequest,
+				ClientMsg: fmt.Sprintf("Invalid alias: %s", err),
+				Err:       err,
+			}
+		}
+
+		alias = &a
+		if t != nil {
+			err = id.ValidateCreateTag(*t)
+			if err != nil {
+				return nil, &api.APIError{
+					Code:      http.StatusBadRequest,
+					ClientMsg: fmt.Sprintf("Invalid tag: %s", err),
+					Err:       err,
+				}
+			}
+
+			tags = []string{*t}
+		}
+	}
+
 	// Create the build
 	data := template.RegisterBuildData{
 		ClusterID:          utils.WithClusterFallback(team.ClusterID),
 		TemplateID:         templateID,
 		UserID:             &userID,
 		Team:               team,
 		Dockerfile:         body.Dockerfile,
-		Alias:              body.Alias,
+		Alias:              alias,
+		Tags:               tags,
 		StartCmd:           body.StartCmd,
 		ReadyCmd:           body.ReadyCmd,
 		CpuCount:           body.CpuCount,

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

@@ -23,7 +23,12 @@ func (a *APIStore) PostV2Templates(c *gin.Context) {
 		return
 	}
 
-	t := requestTemplateBuild(ctx, c, a, api.TemplateBuildRequestV3(body))
+	t := requestTemplateBuild(ctx, c, a, api.TemplateBuildRequestV3{
+		Names:    &[]string{body.Alias},
+		CpuCount: body.CpuCount,
+		MemoryMB: body.MemoryMB,
+		TeamID:   body.TeamID,
+	})
 	if t != nil {
 		c.JSON(http.StatusAccepted, &api.TemplateLegacy{
 			TemplateID: t.TemplateID,

packages/api/internal/db/snapshots_test.go

@@ -71,23 +71,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/deprecated_template_request_build.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.InvalidateAllTags(snapshot.TemplateID)
 
 	return nil
 }