feat: add credits migration script to migrate legacy key credits to credits table

apps/docs/docs.json

@@ -79,9 +79,7 @@
                   },
                   {
                     "group": "Identities",
-                    "pages": [
-                      "quickstart/identities/shared-ratelimits"
-                    ]
+                    "pages": ["quickstart/identities/shared-ratelimits"]
                   }
                 ]
               },
@@ -115,10 +113,7 @@
               {
                 "group": "Identities",
                 "icon": "fingerprint",
-                "pages": [
-                  "concepts/identities/overview",
-                  "concepts/identities/ratelimits"
-                ]
+                "pages": ["concepts/identities/overview", "concepts/identities/ratelimits"]
               }
             ]
           },
@@ -135,9 +130,7 @@
                     "pages": [
                       {
                         "group": "Ratelimiting",
-                        "pages": [
-                          "apis/features/ratelimiting/overview"
-                        ]
+                        "pages": ["apis/features/ratelimiting/overview"]
                       },
                       "apis/features/temp-keys",
                       "apis/features/remaining",
@@ -172,10 +165,7 @@
               {
                 "group": "Audit logs",
                 "icon": "scroll",
-                "pages": [
-                  "audit-log/introduction",
-                  "audit-log/types"
-                ]
+                "pages": ["audit-log/introduction", "audit-log/types"]
               },
               {
                 "group": "Analytics",
@@ -197,10 +187,7 @@
               {
                 "group": "Migrating API Keys",
                 "icon": "plane",
-                "pages": [
-                  "migrations/introduction",
-                  "migrations/keys"
-                ]
+                "pages": ["migrations/introduction", "migrations/keys"]
               }
             ]
           },
@@ -282,9 +269,7 @@
                   },
                   {
                     "group": "Too Many Requests",
-                    "pages": [
-                      "errors/user/too_many_requests/query_quota_exceeded"
-                    ]
+                    "pages": ["errors/user/too_many_requests/query_quota_exceeded"]
                   },
                   {
                     "group": "Unprocessable Entity",
@@ -390,15 +375,11 @@
               },
               {
                 "group": "Go",
-                "pages": [
-                  "libraries/go/api"
-                ]
+                "pages": ["libraries/go/api"]
               },
               {
                 "group": "Python",
-                "pages": [
-                  "libraries/py/api"
-                ]
+                "pages": ["libraries/py/api"]
               }
             ]
           },
@@ -423,15 +404,11 @@
               },
               {
                 "group": "Nuxt",
-                "pages": [
-                  "libraries/nuxt/overview"
-                ]
+                "pages": ["libraries/nuxt/overview"]
               },
               {
                 "group": "Rust",
-                "pages": [
-                  "libraries/rs/overview"
-                ]
+                "pages": ["libraries/rs/overview"]
               },
               {
                 "group": "Springboot",
@@ -474,4 +451,4 @@
     "codeblocks": "system"
   },
   "theme": "maple"
-}
+}

go/cmd/migrate/actions/credits.go

@@ -0,0 +1,153 @@
+package actions
+
+import (
+	"context"
+	"database/sql"
+	"time"
+
+	"github.com/unkeyed/unkey/go/pkg/cli"
+	"github.com/unkeyed/unkey/go/pkg/db"
+	"github.com/unkeyed/unkey/go/pkg/fault"
+	"github.com/unkeyed/unkey/go/pkg/otel/logging"
+	"github.com/unkeyed/unkey/go/pkg/uid"
+)
+
+var CreditsCmd = &cli.Command{
+	Name:  "credits",
+	Usage: "Migrate key credits to separate credits table",
+	Description: `Migrate credit data from the keys table to the separate credits table.
+
+This migration will:
+1. Find all keys with remaining_requests set that don't have credits
+2. Create corresponding entries in the credits table in batches
+3. Preserve refill settings (refill_day, refill_amount, last_refill_at)
+
+The migration is idempotent and will skip keys that already have credits entries.
+
+EXAMPLES:
+unkey migrate credits                    # Run credits migration
+unkey migrate credits --dry-run          # Preview migration without applying changes
+unkey migrate credits --batch-size 1000  # Run migration with custom batch size`,
+	Flags: []cli.Flag{
+		cli.Bool("dry-run", "Preview migration without making changes"),
+		cli.Int("batch-size", "Number of keys to process in each batch", cli.Default(1000)),
+		cli.String("database-primary", "MySQL connection string for primary database. Required for all deployments. Example: user:pass@host:3306/unkey?parseTime=true",
+			cli.Required(), cli.EnvVar("UNKEY_DATABASE_PRIMARY")),
+	},
+	Action: migrateCredits,
+}
+
+func migrateCredits(ctx context.Context, cmd *cli.Command) error {
+	logger := logging.New()
+
+	// Parse flags
+	dryRun := cmd.Bool("dry-run")
+	batchSize := cmd.Int("batch-size")
+	primaryDSN := cmd.String("database-primary")
+
+	if dryRun {
+		logger.Info("Running in dry-run mode - no changes will be made")
+	}
+
+	// Initialize database
+	database, err := db.New(db.Config{
+		PrimaryDSN:  primaryDSN,
+		ReadOnlyDSN: "",
+		Logger:      logger,
+	})
+	if err != nil {
+		return fault.Wrap(err, fault.Internal("Failed to initialize database"))
+	}
+	defer database.Close()
+
+	// Start migration
+	logger.Info("Starting credits migration",
+		"batch_size", batchSize,
+		"dry_run", dryRun,
+	)
+
+	var totalMigrated int
+	offset := 0
+
+	for {
+		// Use the generated sqlc query to find keys without credits
+		keys, err := db.Query.FindKeysWithoutCredits(ctx, database.RO(), db.FindKeysWithoutCreditsParams{
+			Limit:  int32(batchSize),
+			Offset: int32(offset),
+		})
+		if err != nil {
+			return fault.Wrap(err, fault.Internal("Failed to fetch keys"))
+		}
+
+		if len(keys) == 0 {
+			break // No more keys to process
+		}
+
+		logger.Info("Processing batch",
+			"batch_size", len(keys),
+			"offset", offset,
+		)
+
+		if !dryRun {
+			// Prepare batch insert parameters
+			creditParams := make([]db.InsertCreditParams, 0, len(keys))
+
+			for _, key := range keys {
+				creditID := uid.New("credit")
+				now := time.Now().UnixMilli()
+
+				// Convert nullable values
+				remaining := int32(0)
+				if key.RemainingRequests.Valid {
+					remaining = key.RemainingRequests.Int32
+				}
+
+				var refilledAt sql.NullInt64
+				if key.LastRefillAt.Valid {
+					refilledAt = sql.NullInt64{Int64: key.LastRefillAt.Time.UnixMilli(), Valid: true}
+				}
+
+				creditParams = append(creditParams, db.InsertCreditParams{
+					ID:           creditID,
+					WorkspaceID:  key.WorkspaceID,
+					KeyID:        sql.NullString{String: key.ID, Valid: true},
+					IdentityID:   sql.NullString{String: "", Valid: false}, // null for key credits
+					Remaining:    remaining,
+					RefillDay:    key.RefillDay,
+					RefillAmount: key.RefillAmount,
+					CreatedAt:    key.CreatedAtM,
+					UpdatedAt:    sql.NullInt64{Int64: now, Valid: true},
+					RefilledAt:   refilledAt,
+				})
+			}
+
+			err = db.BulkQuery.InsertCredits(ctx, database.RW(), creditParams)
+			if err != nil {
+				logger.Error("Failed to bulk insert credits",
+					"batch_size", len(creditParams),
+					"error", err.Error(),
+				)
+			} else {
+				totalMigrated += len(creditParams)
+			}
+		} else {
+			// Dry run - just count
+			totalMigrated += len(keys)
+		}
+
+		if totalMigrated%10000 == 0 && totalMigrated > 0 {
+			logger.Info("Migration progress",
+				"total_migrated", totalMigrated,
+			)
+		}
+
+		offset += batchSize
+	}
+
+	logger.Info("Migration completed",
+		"total_migrated", totalMigrated,
+		"dry_run", dryRun,
+	)
+
+	return nil
+}

go/cmd/migrate/main.go

@@ -0,0 +1,26 @@
+package migrate
+
+import (
+	"github.com/unkeyed/unkey/go/cmd/migrate/actions"
+	"github.com/unkeyed/unkey/go/pkg/cli"
+)
+
+var Cmd = &cli.Command{
+	Name:  "migrate",
+	Usage: "Run database migrations",
+	Description: `Run various database migrations for Unkey.
+
+This command provides utilities for migrating data between database schemas,
+handling data transformations, and managing database updates.
+
+AVAILABLE MIGRATIONS:
+- credits: Migrate key credits from keys table to separate credits table
+
+EXAMPLES:
+unkey migrate credits                    # Run credits migration
+unkey migrate credits --dry-run          # Preview migration without applying changes
+unkey migrate credits --batch-size 1000  # Run migration with custom batch size`,
+	Commands: []*cli.Command{
+		actions.CreditsCmd,
+	},
+}

go/main.go

@@ -10,6 +10,7 @@ import (
 	dev "github.com/unkeyed/unkey/go/cmd/dev"
 	gateway "github.com/unkeyed/unkey/go/cmd/gw"
 	"github.com/unkeyed/unkey/go/cmd/healthcheck"
+	"github.com/unkeyed/unkey/go/cmd/migrate"
 	"github.com/unkeyed/unkey/go/cmd/quotacheck"
 	"github.com/unkeyed/unkey/go/cmd/run"
 	"github.com/unkeyed/unkey/go/cmd/version"
@@ -35,6 +36,7 @@ func main() {
 			gateway.Cmd,
 			clickhouseUser.Cmd,
 			dev.Cmd,
+			migrate.Cmd,
 		},
 	}