feat(dashboard-api): management route for the projected billing profile

packages/dashboard-api/internal/handlers/management_project_billing_profile.go

@@ -0,0 +1,51 @@
+package handlers
+
+import (
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"go.opentelemetry.io/otel/attribute"
+
+	"github.com/e2b-dev/infra/packages/dashboard-api/internal/api"
+	"github.com/e2b-dev/infra/packages/dashboard-api/internal/management"
+	"github.com/e2b-dev/infra/packages/shared/pkg/ginutils"
+	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
+
+// ManagementUpsertProjectBillingProfile records a project's billing profile as
+// resolved on the global billing plane.
+//
+// Idempotent for the same reason the limits route is: the revision decides
+// which delivery writes, and a stale delivery is answered 204 — what it asked
+// for is stored, and a newer answer is not this call's to undo.
+func (s *APIStore) ManagementUpsertProjectBillingProfile(c *gin.Context, projectID api.ProjectID) {
+	ctx := c.Request.Context()
+	attrs := []attribute.KeyValue{telemetry.WithTeamID(projectID.String())}
+
+	body, err := ginutils.ParseBody[api.ManagementProjectBillingProfile](ctx, c)
+	if err != nil {
+		telemetry.ReportErrorByCode(ctx, http.StatusBadRequest, "upsert project billing profile failed",
+			fmt.Errorf("parse project billing profile request: %w", err), attrs...)
+		s.sendAPIStoreError(c, http.StatusBadRequest, "Invalid request body")
+
+		return
+	}
+
+	if err := s.managementService.ApplyBillingProfile(ctx, management.BillingProfileProjection{
+		ProjectID:        projectID,
+		Revision:         body.Revision,
+		DecidedAt:        utils.DerefOrDefault(body.DecidedAt, time.Time{}),
+		HasPaymentMethod: body.HasPaymentMethod,
+		Enterprise:       body.Enterprise,
+		Plan:             utils.DerefOrDefault(body.Plan, ""),
+	}); err != nil {
+		s.sendBillingProfileError(c, err, attrs...)
+
+		return
+	}
+
+	c.Status(http.StatusNoContent)
+}

packages/dashboard-api/internal/handlers/utils_management.go

@@ -57,6 +57,22 @@ func (s *APIStore) sendProjectBlockError(c *gin.Context, err error, attrs ...att
 	}
 }
 
+func (s *APIStore) sendBillingProfileError(c *gin.Context, err error, attrs ...attribute.KeyValue) {
+	ctx := c.Request.Context()
+
+	switch {
+	case errors.Is(err, management.ErrProjectNotFound):
+		telemetry.ReportErrorByCode(ctx, http.StatusNotFound, "apply billing profile failed", err, attrs...)
+		s.sendAPIStoreError(c, http.StatusNotFound, "Project not found")
+	case errors.Is(err, management.ErrInvalidBillingProfile):
+		telemetry.ReportErrorByCode(ctx, http.StatusBadRequest, "apply billing profile failed", err, attrs...)
+		s.sendAPIStoreError(c, http.StatusBadRequest, "Invalid billing profile")
+	default:
+		telemetry.ReportCriticalError(ctx, "apply billing profile failed", err, attrs...)
+		s.sendAPIStoreError(c, http.StatusInternalServerError, "Error applying billing profile")
+	}
+}
+
 // A cache eviction that failed is reported: the values are committed, so the
 // caller retrying is what gets the stale entry cleared.
 func (s *APIStore) sendProjectLimitsError(c *gin.Context, err error, attrs ...attribute.KeyValue) {

packages/dashboard-api/internal/management/apply_lag.go

@@ -20,8 +20,9 @@ const (
 type projectionKind string
 
 const (
-	projectionProjectLimits projectionKind = "project_limits"
-	projectionProjectBlocks projectionKind = "project_blocks"
+	projectionProjectLimits   projectionKind = "project_limits"
+	projectionProjectBlocks   projectionKind = "project_blocks"
+	projectionBillingProfiles projectionKind = "billing_profiles"
 )
 
 type applyLag struct {

packages/dashboard-api/internal/management/billing_profile.go

@@ -0,0 +1,117 @@
+package management
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/google/uuid"
+	"go.uber.org/zap"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/dberrors"
+	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+)
+
+var ErrInvalidBillingProfile = errors.New("invalid billing profile")
+
+// The facts arrive resolved: the caller owns the global billing state and this
+// side stores what it is given. A team without a row has an unknown profile —
+// readers must not default it.
+type BillingProfileProjection struct {
+	ProjectID uuid.UUID
+	Revision  int64
+	// Zero when the caller did not say.
+	DecidedAt time.Time
+
+	HasPaymentMethod bool
+	Enterprise       bool
+	// Empty when no plan is in force or the caller does not know.
+	Plan string
+}
+
+// ApplyBillingProfile records a project's billing profile, behind the revision
+// that decides whether this delivery is the newest one.
+func (s *Service) ApplyBillingProfile(ctx context.Context, projection BillingProfileProjection) error {
+	if projection.ProjectID == uuid.Nil || projection.Revision <= 0 {
+		return ErrInvalidBillingProfile
+	}
+
+	stored, err := s.applyBillingProfile(ctx, projection)
+	if err != nil {
+		return err
+	}
+	if stored {
+		s.applyLag.stored(ctx, projectionBillingProfiles, projection.DecidedAt)
+	}
+
+	// Logged rather than returned, as the sibling writes do: the row is
+	// committed, and a caller told the delivery failed would repeat a revision
+	// this side already has.
+	if err := s.cache.InvalidateTeamCache(ctx, projection.ProjectID); err != nil {
+		logger.L().Error(ctx, "invalidating team cache after billing profile update",
+			logger.WithTeamID(projection.ProjectID.String()), zap.Error(err))
+	}
+
+	return nil
+}
+
+func (s *Service) applyBillingProfile(ctx context.Context, projection BillingProfileProjection) (bool, error) {
+	txDB, tx, err := s.projectDB.WithTx(ctx)
+	if err != nil {
+		return false, fmt.Errorf("start billing profile transaction: %w", err)
+	}
+	defer func() {
+		_ = tx.Rollback(ctx)
+	}()
+
+	if _, err := txDB.LockManagedProject(ctx, projection.ProjectID); err != nil {
+		if dberrors.IsNotFoundError(err) {
+			return false, ErrProjectNotFound
+		}
+
+		return false, fmt.Errorf("lock project: %w", err)
+	}
+
+	applied, err := txDB.ApplyBillingProfileProjection(ctx, queries.ApplyBillingProfileProjectionParams{
+		ProjectID: projection.ProjectID,
+		Revision:  projection.Revision,
+		DecidedAt: decidedAtParam(projection.DecidedAt),
+	})
+	if err != nil {
+		return false, fmt.Errorf("advance billing profile projection: %w", err)
+	}
+	if !applied {
+		if err := tx.Commit(ctx); err != nil {
+			return false, fmt.Errorf("commit stale billing profile projection: %w", err)
+		}
+
+		return false, nil
+	}
+
+	if err := txDB.UpsertTeamBillingProfile(ctx, queries.UpsertTeamBillingProfileParams{
+		TeamID:           projection.ProjectID,
+		HasPaymentMethod: projection.HasPaymentMethod,
+		Enterprise:       projection.Enterprise,
+		Plan:             planParam(projection.Plan),
+	}); err != nil {
+		return false, fmt.Errorf("upsert team billing profile: %w", err)
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return false, fmt.Errorf("commit billing profile projection: %w", err)
+	}
+
+	return true, nil
+}
+
+func planParam(plan string) *string {
+	trimmed := strings.TrimSpace(plan)
+	if trimmed == "" {
+		return nil
+	}
+
+	return &trimmed
+}

packages/dashboard-api/internal/management/billing_profile_test.go

@@ -0,0 +1,121 @@
+package management
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/metric/noop"
+
+	"github.com/e2b-dev/infra/packages/db/pkg/testutils"
+)
+
+func TestApplyBillingProfileWritesAndAdvances(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	service := NewService(db.AuthDB, db.SqlcClient, &countingProfileCache{}, noop.NewMeterProvider())
+	teamID := testutils.CreateTestTeam(t, db)
+
+	require.NoError(t, service.ApplyBillingProfile(t.Context(), BillingProfileProjection{
+		ProjectID: teamID, Revision: 1, HasPaymentMethod: false, Enterprise: false,
+	}))
+	hasPM, enterprise, plan := billingProfileState(t, db, teamID)
+	require.False(t, hasPM)
+	require.False(t, enterprise)
+	require.Nil(t, plan)
+
+	require.NoError(t, service.ApplyBillingProfile(t.Context(), BillingProfileProjection{
+		ProjectID: teamID, Revision: 2, HasPaymentMethod: true, Enterprise: false, Plan: "pro_v1",
+	}))
+	hasPM, _, plan = billingProfileState(t, db, teamID)
+	require.True(t, hasPM)
+	require.NotNil(t, plan)
+	require.Equal(t, "pro_v1", *plan)
+}
+
+func TestApplyBillingProfileDropsStaleRevision(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	service := NewService(db.AuthDB, db.SqlcClient, &countingProfileCache{}, noop.NewMeterProvider())
+	teamID := testutils.CreateTestTeam(t, db)
+
+	require.NoError(t, service.ApplyBillingProfile(t.Context(), BillingProfileProjection{
+		ProjectID: teamID, Revision: 5, HasPaymentMethod: true, Enterprise: true,
+	}))
+	// A delayed duplicate and an older delivery both succeed without writing.
+	for _, revision := range []int64{5, 4} {
+		require.NoError(t, service.ApplyBillingProfile(t.Context(), BillingProfileProjection{
+			ProjectID: teamID, Revision: revision, HasPaymentMethod: false, Enterprise: false,
+		}))
+	}
+	hasPM, enterprise, _ := billingProfileState(t, db, teamID)
+	require.True(t, hasPM)
+	require.True(t, enterprise)
+}
+
+func TestApplyBillingProfileRejectsInvalidAndMissing(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	service := NewService(db.AuthDB, db.SqlcClient, &countingProfileCache{}, noop.NewMeterProvider())
+
+	require.ErrorIs(t, service.ApplyBillingProfile(t.Context(),
+		BillingProfileProjection{ProjectID: uuid.Nil, Revision: 1}), ErrInvalidBillingProfile)
+	require.ErrorIs(t, service.ApplyBillingProfile(t.Context(),
+		BillingProfileProjection{ProjectID: uuid.New(), Revision: 0}), ErrInvalidBillingProfile)
+	require.ErrorIs(t, service.ApplyBillingProfile(t.Context(),
+		BillingProfileProjection{ProjectID: uuid.New(), Revision: 1}), ErrProjectNotFound)
+}
+
+func TestApplyBillingProfileSurvivesFailedCacheInvalidation(t *testing.T) {
+	t.Parallel()
+
+	db := testutils.SetupDatabase(t)
+	cache := &countingProfileCache{}
+	service := NewService(db.AuthDB, db.SqlcClient, cache, noop.NewMeterProvider())
+	teamID := testutils.CreateTestTeam(t, db)
+
+	// Unlike blocks, a failed eviction is logged, not returned: the row is
+	// committed and the caller must not re-send a revision this side has.
+	cache.err = errors.New("cache unavailable")
+	require.NoError(t, service.ApplyBillingProfile(t.Context(), BillingProfileProjection{
+		ProjectID: teamID, Revision: 1, HasPaymentMethod: true, Enterprise: false,
+	}))
+	hasPM, _, _ := billingProfileState(t, db, teamID)
+	require.True(t, hasPM)
+	require.Equal(t, 1, cache.calls)
+}
+
+type countingProfileCache struct {
+	noopAuthService
+
+	err   error
+	calls int
+}
+
+func (c *countingProfileCache) InvalidateTeamCache(context.Context, uuid.UUID) error {
+	c.calls++
+
+	return c.err
+}
+
+func billingProfileState(t *testing.T, db *testutils.Database, teamID uuid.UUID) (bool, bool, *string) {
+	t.Helper()
+
+	var hasPM, enterprise bool
+	var plan *string
+	require.NoError(t, db.SqlcClient.TestsRawSQLQuery(t.Context(),
+		"SELECT has_payment_method, enterprise, plan FROM public.team_billing_profiles WHERE team_id = $1",
+		func(rows pgx.Rows) error {
+			rows.Next()
+
+			return rows.Scan(&hasPM, &enterprise, &plan)
+		}, teamID))
+
+	return hasPM, enterprise, plan
+}

packages/db/migrations/20261007142026_add_team_billing_profiles.sql

@@ -0,0 +1,40 @@
+-- +goose Up
+-- +goose StatementBegin
+
+SET LOCAL lock_timeout = '5s';
+
+-- Delivery ledger for the pushed billing profile, same shape and reasoning as
+-- projection.project_limits: the push is at-least-once over a network, so the
+-- older of two in-flight deliveries has to be refused where it lands, in the
+-- same transaction as the values it would have written.
+CREATE TABLE projection.billing_profiles (
+    project_id uuid PRIMARY KEY REFERENCES public.teams(id) ON DELETE CASCADE,
+    revision bigint NOT NULL CHECK (revision > 0),
+    decided_at timestamptz,
+    created_at timestamptz NOT NULL DEFAULT now(),
+    updated_at timestamptz NOT NULL DEFAULT now()
+);
+
+-- The billing facts a plane cannot know on its own once a team's billing lives
+-- on the global plane: whether the workspace can be charged and which plan is
+-- in force. Regional readers (bot-detection's payment signal, the LaunchDarkly
+-- tier context, tier labels) prefer this row over the legacy regional sources;
+-- a team without a row has an unknown profile, never a defaulted one.
+CREATE TABLE public.team_billing_profiles (
+    team_id uuid PRIMARY KEY REFERENCES public.teams(id) ON DELETE CASCADE,
+    has_payment_method boolean NOT NULL,
+    enterprise boolean NOT NULL,
+    plan text,
+    updated_at timestamptz NOT NULL DEFAULT now()
+);
+
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+
+SET LOCAL lock_timeout = '5s';
+DROP TABLE public.team_billing_profiles;
+DROP TABLE projection.billing_profiles;
+
+-- +goose StatementEnd