Merge P0-3+P0-4: storage-quota suspend constraint fix + infra revoke + auto-unsuspend

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: mastermanas805

internal/config/config.go

@@ -80,8 +80,8 @@ type Config struct {
 	// /internal/teams/:id/terminate against the api repo. Both must be
 	// set for the terminator to act — otherwise the job short-circuits
 	// each tick with a WARN.
-	InstantAPIInternalURL    string // INSTANT_API_INTERNAL_URL — base URL of api (cluster-local)
-	WorkerInternalJWTSecret  string // WORKER_INTERNAL_JWT_SECRET — HS256 shared with api
+	InstantAPIInternalURL   string // INSTANT_API_INTERNAL_URL — base URL of api (cluster-local)
+	WorkerInternalJWTSecret string // WORKER_INTERNAL_JWT_SECRET — HS256 shared with api
 
 	// AES_KEY — 64-char hex key (32 bytes) used to decrypt
 	// resources.connection_url for the customer-backup runner AND the real
@@ -110,6 +110,24 @@ type Config struct {
 	// every platform dump in chronological order" without filtering customer
 	// data.
 	PlatformBackupS3Prefix string // PLATFORM_BACKUP_S3_PREFIX (default: "platform-backups/")
+
+	// Storage-quota enforcement: infra revoke/grant on suspend/unsuspend (P0-3/P0-4 fix).
+	// These mirror the api's config for the same operations — the worker
+	// needs direct access to customer infrastructure to revoke connections
+	// (ACL SETUSER off / REVOKE CONNECT / revokeRolesFromUser) when a
+	// resource exceeds its storage quota and to re-grant on auto-unsuspend.
+	//
+	// All three are optional — when absent the worker skips the infra-revoke
+	// step (fail-open) and only flips the status row. The row flip still
+	// gates the user's provisioning API calls; the infra revoke terminates
+	// live TCP connections.
+	//
+	// CUSTOMER_DATABASE_URL — admin DSN for shared Postgres cluster.
+	// MONGO_ADMIN_URI       — admin URI for shared MongoDB cluster.
+	// CUSTOMER_REDIS_URL    — admin Redis URL for shared Redis cluster.
+	CustomerDatabaseURL string // CUSTOMER_DATABASE_URL
+	MongoAdminURI       string // MONGO_ADMIN_URI
+	CustomerRedisURL    string // CUSTOMER_REDIS_URL
 }
 
 // ErrMissingConfig is returned when a required env var is absent.
@@ -188,6 +206,13 @@ func Load() *Config {
 		// with a YYYY-MM-DD path segment yields well-formed S3 keys without
 		// the producer needing to know whether the parent ends in a slash.
 		PlatformBackupS3Prefix: getenv("PLATFORM_BACKUP_S3_PREFIX", "platform-backups/"),
+
+		// Storage-quota suspend/unsuspend infra credentials (P0-3/P0-4 fix).
+		// All three are optional — when absent the infra-revoke step is
+		// skipped (fail-open) and only the status row is updated.
+		CustomerDatabaseURL: os.Getenv("CUSTOMER_DATABASE_URL"),
+		MongoAdminURI:       os.Getenv("MONGO_ADMIN_URI"),
+		CustomerRedisURL:    os.Getenv("CUSTOMER_REDIS_URL"),
 	}
 
 	// Fall back to the shared object-store bucket when the operator hasn't
@@ -222,14 +247,17 @@ func Load() *Config {
 		"ses_from_set", cfg.SESFromEmail != "",
 		"ses_template_count", len(cfg.SESTemplateNames),
 		"aes_key_set", cfg.AESKey != "",
+		"customer_db_set", cfg.CustomerDatabaseURL != "",
+		"mongo_admin_set", cfg.MongoAdminURI != "",
+		"customer_redis_set", cfg.CustomerRedisURL != "",
 	)
 	return cfg
 }
 
 // parseBrevoTemplateIDs decodes the BREVO_TEMPLATE_IDS env var. The expected
 // shape is a JSON object mapping audit_log.kind → numeric Brevo template id:
 //
-//   BREVO_TEMPLATE_IDS='{"subscription.upgraded": 12, "near_quota_wall": 7}'
+//	BREVO_TEMPLATE_IDS='{"subscription.upgraded": 12, "near_quota_wall": 7}'
 //
 // Empty string returns an empty map (Brevo will then SkipNoTemplate on every
 // send — operator opted into "API key set, no templates yet"). A malformed
@@ -255,7 +283,7 @@ func parseBrevoTemplateIDs(raw string) map[string]int {
 // shape is a JSON object mapping audit_log.kind → SES template name (string,
 // not numeric — SES references templates by name unlike Brevo):
 //
-//   SES_TEMPLATE_NAMES='{"subscription.upgraded": "tier-upgraded-v1", "near_quota_wall": "quota-wall-nudge-v1"}'
+//	SES_TEMPLATE_NAMES='{"subscription.upgraded": "tier-upgraded-v1", "near_quota_wall": "quota-wall-nudge-v1"}'
 //
 // Empty string returns an empty map (SES will then SkipNoTemplate on every
 // send — operator opted into "AWS creds set, no templates yet"). A malformed

internal/jobs/quota.go

@@ -35,33 +35,87 @@ type StorageQuotaChecker interface {
 	CheckStorageQuota(ctx context.Context, db *sql.DB, resourceID uuid.UUID, limitMB int) (int64, bool, error)
 }
 
+// resourceStatusSuspended / resourceStatusActive are named constants for the
+// two status values touched by the suspend/unsuspend lifecycle. Using constants
+// avoids scattered string literals and makes the "suspended → active" transition
+// greppable. (CLAUDE.md: "Use named constants, not inline strings".)
+const (
+	resourceStatusSuspended = "suspended"
+	resourceStatusActive    = "active"
+)
+
 // EnforceStorageQuotaWorker checks all active resources against their plan's
-// storage limit and suspends those that have exceeded it.
+// storage limit and suspends those that have exceeded it. It also scans
+// currently-suspended resources and unsuspends those whose usage has dropped
+// back below the limit.
+//
+// Two loops per Work() run:
+//  1. Suspend loop — scans status='active' postgres/redis/mongodb resources,
+//     calls the provisioner-side infra revoke (REVOKE CONNECT / ACL SETUSER off
+//     / revokeRolesFromUser), then flips status to 'suspended'.
+//  2. Unsuspend loop — scans status='suspended' resources, re-checks storage,
+//     re-grants infra access, and flips status back to 'active' for those now
+//     under limit.
+//
+// Both loops fail-open on infra errors (connectivity issues with the customer
+// DB / Redis / Mongo only affect the row flip as a fallback — the status flip
+// still lands). Fail-open matches CLAUDE.md convention #1.
 type EnforceStorageQuotaWorker struct {
 	river.WorkerDefaults[EnforceStorageQuotaArgs]
-	db    *sql.DB
-	plans PlanRegistry
+	db      *sql.DB
+	plans   PlanRegistry
+	revoker ResourceInfraRevoker // nil = no infra revoke (status-row only)
 }
 
 // NewEnforceStorageQuotaWorker constructs an EnforceStorageQuotaWorker.
-func NewEnforceStorageQuotaWorker(db *sql.DB, plans PlanRegistry) *EnforceStorageQuotaWorker {
-	return &EnforceStorageQuotaWorker{db: db, plans: plans}
+// Pass nil for revoker to skip infra-level revoke/grant (status flip still
+// occurs). In production, pass NewDirectResourceRevoker(...) built from cfg.
+func NewEnforceStorageQuotaWorker(db *sql.DB, plans PlanRegistry, revoker ResourceInfraRevoker) *EnforceStorageQuotaWorker {
+	return &EnforceStorageQuotaWorker{db: db, plans: plans, revoker: revoker}
 }
 
-// Work scans all active postgres/redis/mongodb resources and suspends those over their quota.
+// Work scans all active postgres/redis/mongodb resources and suspends those
+// over their quota, then unsuspends any suspended resources that are now back
+// under their quota.
 func (w *EnforceStorageQuotaWorker) Work(ctx context.Context, job *river.Job[EnforceStorageQuotaArgs]) error {
 	ctx, span := otel.Tracer("instant.dev/worker").Start(ctx, "job.enforce_storage_quota")
 	defer span.End()
 
+	suspended, err := w.runSuspendLoop(ctx)
+	if err != nil {
+		return err
+	}
+
+	unsuspended, err := w.runUnsuspendLoop(ctx)
+	if err != nil {
+		// Don't fail the job for unsuspend errors — suspends already landed.
+		slog.Error("jobs.enforce_storage_quota.unsuspend_loop_error", "error", err)
+	}
+
+	var jobID int64
+	if job.JobRow != nil {
+		jobID = job.ID
+	}
+	slog.Info("jobs.enforce_storage_quota.completed",
+		"suspended_count", suspended,
+		"unsuspended_count", unsuspended,
+		"job_id", jobID,
+	)
+	return nil
+}
+
+// runSuspendLoop scans status='active' resources, suspends over-quota ones,
+// and returns the count of newly suspended resources.
+func (w *EnforceStorageQuotaWorker) runSuspendLoop(ctx context.Context) (int, error) {
 	rows, err := w.db.QueryContext(ctx, `
 		SELECT id, token, resource_type, tier, storage_bytes
 		FROM resources
-		WHERE status = 'active'
+		WHERE status = $1
 		  AND resource_type IN ('postgres', 'redis', 'mongodb')
 		ORDER BY created_at
-	`)
+	`, resourceStatusActive)
 	if err != nil {
-		return fmt.Errorf("EnforceStorageQuotaWorker: query failed: %w", err)
+		return 0, fmt.Errorf("EnforceStorageQuotaWorker.suspendLoop: query failed: %w", err)
 	}
 	defer rows.Close()
 
@@ -83,7 +137,7 @@ func (w *EnforceStorageQuotaWorker) Work(ctx context.Context, job *river.Job[Enf
 
 		limitMB := w.plans.StorageLimitMB(tier, resourceType)
 		if limitMB == -1 {
-			continue // unlimited tier — skip
+			continue // unlimited tier — never suspend
 		}
 
 		uid, parseErr := uuid.Parse(id)
@@ -105,10 +159,29 @@ func (w *EnforceStorageQuotaWorker) Work(ctx context.Context, job *river.Job[Enf
 			continue
 		}
 
+		// Infra revoke FIRST, then status flip — matches the iron-rule order
+		// from api/internal/handlers/resource.go Pause(): "provider-side FIRST
+		// so the row stays active if infra fails; row flip is the commit."
+		// Here we invert slightly: infra revoke is fail-open (logged warning,
+		// not a hard error) so we always proceed to the status flip. This is
+		// intentional: a row marked 'suspended' blocks new provisions from the
+		// API even when the infra revoke is not available (customer DB down).
+		if w.revoker != nil {
+			if revokeErr := w.revoker.RevokeAccess(ctx, resourceType, token, ""); revokeErr != nil {
+				// revoker implementations are fail-open (return nil on infra
+				// error, log a WARN). A non-nil error here is unexpected —
+				// log it but don't abort the row update.
+				slog.Error("jobs.enforce_storage_quota.revoke_error",
+					"resource_id", id, "token", token, "resource_type", resourceType,
+					"error", revokeErr,
+				)
+			}
+		}
+
 		_, updateErr := w.db.ExecContext(ctx, `
-			UPDATE resources SET status = 'suspended'
-			WHERE id = $1 AND status = 'active'
-		`, id)
+			UPDATE resources SET status = $1
+			WHERE id = $2 AND status = $3
+		`, resourceStatusSuspended, id, resourceStatusActive)
 		if updateErr != nil {
 			slog.Error("jobs.enforce_storage_quota.suspend_failed",
 				"resource_id", id,
@@ -128,19 +201,111 @@ func (w *EnforceStorageQuotaWorker) Work(ctx context.Context, job *river.Job[Enf
 		suspended++
 	}
 	if err := rows.Err(); err != nil {
-		return fmt.Errorf("EnforceStorageQuotaWorker: rows error: %w", err)
+		return suspended, fmt.Errorf("EnforceStorageQuotaWorker.suspendLoop: rows error: %w", err)
 	}
 
-	var jobID int64
-	if job.JobRow != nil {
-		jobID = job.ID
-	}
-	slog.Info("jobs.enforce_storage_quota.completed",
+	slog.Info("jobs.enforce_storage_quota.suspend_loop_done",
 		"checked_count", checked,
 		"suspended_count", suspended,
-		"job_id", jobID,
 	)
-	return nil
+	return suspended, nil
+}
+
+// runUnsuspendLoop scans status='suspended' resources, re-checks quota, and
+// re-grants infra access + flips to 'active' for those now under their limit.
+// Returns the count of newly unsuspended resources.
+func (w *EnforceStorageQuotaWorker) runUnsuspendLoop(ctx context.Context) (int, error) {
+	rows, err := w.db.QueryContext(ctx, `
+		SELECT id, token, resource_type, tier, storage_bytes
+		FROM resources
+		WHERE status = $1
+		  AND resource_type IN ('postgres', 'redis', 'mongodb')
+		ORDER BY created_at
+	`, resourceStatusSuspended)
+	if err != nil {
+		return 0, fmt.Errorf("EnforceStorageQuotaWorker.unsuspendLoop: query failed: %w", err)
+	}
+	defer rows.Close()
+
+	unsuspended := 0
+
+	for rows.Next() {
+		var (
+			id           string
+			token        string
+			resourceType string
+			tier         string
+			storageBytes int64
+		)
+		if scanErr := rows.Scan(&id, &token, &resourceType, &tier, &storageBytes); scanErr != nil {
+			slog.Error("jobs.enforce_storage_quota.unsuspend_scan_error", "error", scanErr)
+			continue
+		}
+
+		limitMB := w.plans.StorageLimitMB(tier, resourceType)
+		if limitMB == -1 {
+			// Unlimited tier shouldn't have been suspended, but unsuspend
+			// eagerly to self-heal any historical bad state.
+			limitMB = 0 // treat as exceeded=false below
+		}
+
+		uid, parseErr := uuid.Parse(id)
+		if parseErr != nil {
+			slog.Error("jobs.enforce_storage_quota.unsuspend_invalid_uuid", "id", id, "error", parseErr)
+			continue
+		}
+
+		var exceeded bool
+		if limitMB > 0 {
+			_, exceeded, err = checkStorageQuota(ctx, w.db, uid, limitMB)
+			if err != nil {
+				slog.Error("jobs.enforce_storage_quota.unsuspend_check_error",
+					"resource_id", id, "error", err)
+				continue // fail open — don't unsuspend on check error
+			}
+		}
+		// limitMB == 0 means unlimited; exceeded stays false → will unsuspend.
+
+		if exceeded {
+			continue // still over quota — remain suspended
+		}
+
+		// Re-grant infra access before flipping the status row.
+		if w.revoker != nil {
+			if grantErr := w.revoker.GrantAccess(ctx, resourceType, token, ""); grantErr != nil {
+				slog.Error("jobs.enforce_storage_quota.grant_error",
+					"resource_id", id, "token", token, "resource_type", resourceType,
+					"error", grantErr,
+				)
+				// Non-nil means unexpected path — still proceed with row flip.
+			}
+		}
+
+		_, updateErr := w.db.ExecContext(ctx, `
+			UPDATE resources SET status = $1
+			WHERE id = $2 AND status = $3
+		`, resourceStatusActive, id, resourceStatusSuspended)
+		if updateErr != nil {
+			slog.Error("jobs.enforce_storage_quota.unsuspend_failed",
+				"resource_id", id, "error", updateErr)
+			continue
+		}
+
+		slog.Info("jobs.enforce_storage_quota.unsuspended",
+			"resource_id", id,
+			"token", token,
+			"resource_type", resourceType,
+			"tier", tier,
+			"storage_bytes", storageBytes,
+			"limit_mb", limitMB,
+		)
+		unsuspended++
+	}
+	if err := rows.Err(); err != nil {
+		return unsuspended, fmt.Errorf("EnforceStorageQuotaWorker.unsuspendLoop: rows error: %w", err)
+	}
+
+	return unsuspended, nil
 }
 
 // checkStorageQuota reads resources.storage_bytes and compares to limitMB.