diff --git a/apps/docs/docs.json b/apps/docs/docs.json index ed595cb096..0f1c806bd8 100644 --- a/apps/docs/docs.json +++ b/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" -} \ No newline at end of file +} diff --git a/go/cmd/migrate/actions/credits.go b/go/cmd/migrate/actions/credits.go new file mode 100644 index 0000000000..43f6abb115 --- /dev/null +++ b/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 +} diff --git a/go/cmd/migrate/main.go b/go/cmd/migrate/main.go new file mode 100644 index 0000000000..56ea3426ab --- /dev/null +++ b/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, + }, +} diff --git a/go/main.go b/go/main.go index 0f2ec0c5f8..406b96779d 100644 --- a/go/main.go +++ b/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, }, } diff --git a/go/pkg/db/keys_find_without_credits.sql_generated.go b/go/pkg/db/keys_find_without_credits.sql_generated.go index 053ec83e3b..8527dc2897 100644 --- a/go/pkg/db/keys_find_without_credits.sql_generated.go +++ b/go/pkg/db/keys_find_without_credits.sql_generated.go @@ -11,23 +11,20 @@ import ( ) const findKeysWithoutCredits = `-- name: FindKeysWithoutCredits :many -SELECT +SELECT k.id, k.workspace_id, k.remaining_requests, k.refill_day, k.refill_amount, - CASE - WHEN k.last_refill_at IS NULL THEN NULL - ELSE UNIX_TIMESTAMP(k.last_refill_at) * 1000 - END as last_refill_at_unix, - k.created_at_m, - k.updated_at_m + k.last_refill_at, + k.created_at_m FROM ` + "`" + `keys` + "`" + ` k LEFT JOIN ` + "`" + `credits` + "`" + ` c ON c.key_id = k.id -WHERE k.deleted_at_m IS NULL - AND k.remaining_requests IS NOT NULL +LEFT JOIN ` + "`" + `credits` + "`" + ` c2 ON c2.identity_id = k.identity_id +WHERE k.remaining_requests IS NOT NULL AND c.id IS NULL + AND c2.id IS NULL ORDER BY k.created_at_m DESC LIMIT ? OFFSET ? @@ -44,9 +41,8 @@ type FindKeysWithoutCreditsRow struct { RemainingRequests sql.NullInt32 `db:"remaining_requests"` RefillDay sql.NullInt16 `db:"refill_day"` RefillAmount sql.NullInt32 `db:"refill_amount"` - LastRefillAtUnix interface{} `db:"last_refill_at_unix"` + LastRefillAt sql.NullTime `db:"last_refill_at"` CreatedAtM int64 `db:"created_at_m"` - UpdatedAtM sql.NullInt64 `db:"updated_at_m"` } // FindKeysWithoutCredits @@ -57,17 +53,14 @@ type FindKeysWithoutCreditsRow struct { // k.remaining_requests, // k.refill_day, // k.refill_amount, -// CASE -// WHEN k.last_refill_at IS NULL THEN NULL -// ELSE UNIX_TIMESTAMP(k.last_refill_at) * 1000 -// END as last_refill_at_unix, -// k.created_at_m, -// k.updated_at_m +// k.last_refill_at, +// k.created_at_m // FROM `keys` k // LEFT JOIN `credits` c ON c.key_id = k.id -// WHERE k.deleted_at_m IS NULL -// AND k.remaining_requests IS NOT NULL +// LEFT JOIN `credits` c2 ON c2.identity_id = k.identity_id +// WHERE k.remaining_requests IS NOT NULL // AND c.id IS NULL +// AND c2.id IS NULL // ORDER BY k.created_at_m DESC // LIMIT ? // OFFSET ? @@ -86,9 +79,8 @@ func (q *Queries) FindKeysWithoutCredits(ctx context.Context, db DBTX, arg FindK &i.RemainingRequests, &i.RefillDay, &i.RefillAmount, - &i.LastRefillAtUnix, + &i.LastRefillAt, &i.CreatedAtM, - &i.UpdatedAtM, ); err != nil { return nil, err } diff --git a/go/pkg/db/querier_generated.go b/go/pkg/db/querier_generated.go index 0490abba2c..3955e24bb2 100644 --- a/go/pkg/db/querier_generated.go +++ b/go/pkg/db/querier_generated.go @@ -495,17 +495,14 @@ type Querier interface { // k.remaining_requests, // k.refill_day, // k.refill_amount, - // CASE - // WHEN k.last_refill_at IS NULL THEN NULL - // ELSE UNIX_TIMESTAMP(k.last_refill_at) * 1000 - // END as last_refill_at_unix, - // k.created_at_m, - // k.updated_at_m + // k.last_refill_at, + // k.created_at_m // FROM `keys` k // LEFT JOIN `credits` c ON c.key_id = k.id - // WHERE k.deleted_at_m IS NULL - // AND k.remaining_requests IS NOT NULL + // LEFT JOIN `credits` c2 ON c2.identity_id = k.identity_id + // WHERE k.remaining_requests IS NOT NULL // AND c.id IS NULL + // AND c2.id IS NULL // ORDER BY k.created_at_m DESC // LIMIT ? // OFFSET ? diff --git a/go/pkg/db/queries/keys_find_without_credits.sql b/go/pkg/db/queries/keys_find_without_credits.sql index 532916c41d..982da4f9f8 100644 --- a/go/pkg/db/queries/keys_find_without_credits.sql +++ b/go/pkg/db/queries/keys_find_without_credits.sql @@ -1,21 +1,18 @@ -- name: FindKeysWithoutCredits :many -SELECT +SELECT k.id, k.workspace_id, k.remaining_requests, k.refill_day, k.refill_amount, - CASE - WHEN k.last_refill_at IS NULL THEN NULL - ELSE UNIX_TIMESTAMP(k.last_refill_at) * 1000 - END as last_refill_at_unix, - k.created_at_m, - k.updated_at_m + k.last_refill_at, + k.created_at_m FROM `keys` k LEFT JOIN `credits` c ON c.key_id = k.id -WHERE k.deleted_at_m IS NULL - AND k.remaining_requests IS NOT NULL +LEFT JOIN `credits` c2 ON c2.identity_id = k.identity_id +WHERE k.remaining_requests IS NOT NULL AND c.id IS NULL + AND c2.id IS NULL ORDER BY k.created_at_m DESC LIMIT ? -OFFSET ?; \ No newline at end of file +OFFSET ?;