diff --git a/go/apps/api/routes/v2_identities_delete_identity/handler.go b/go/apps/api/routes/v2_identities_delete_identity/handler.go index 404e799cb8..20c5eb65ba 100644 --- a/go/apps/api/routes/v2_identities_delete_identity/handler.go +++ b/go/apps/api/routes/v2_identities_delete_identity/handler.go @@ -71,27 +71,25 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - results, err := db.Query.FindIdentityWithRatelimits(ctx, h.DB.RO(), db.FindIdentityWithRatelimitsParams{ + identity, err := db.Query.FindIdentity(ctx, h.DB.RO(), db.FindIdentityParams{ WorkspaceID: auth.AuthorizedWorkspaceID, Identity: req.Identity, Deleted: false, }) if err != nil { + if db.IsNotFound(err) { + return fault.New("identity not found", + fault.Code(codes.Data.Identity.NotFound.URN()), + fault.Internal("identity not found"), fault.Public("This identity does not exist."), + ) + } + return fault.Wrap(err, fault.Code(codes.App.Internal.ServiceUnavailable.URN()), fault.Internal("database failed to find the identity"), fault.Public("Error finding the identity."), ) } - if len(results) == 0 { - return fault.New("identity not found", - fault.Code(codes.Data.Identity.NotFound.URN()), - fault.Internal("identity not found"), fault.Public("This identity does not exist."), - ) - } - - identity := results[0] - // Parse ratelimits JSON var ratelimits []db.RatelimitInfo if ratelimitBytes, ok := identity.Ratelimits.([]byte); ok && ratelimitBytes != nil { diff --git a/go/apps/api/routes/v2_identities_get_identity/handler.go b/go/apps/api/routes/v2_identities_get_identity/handler.go index 4b1b17fe70..1e0827d0cb 100644 --- a/go/apps/api/routes/v2_identities_get_identity/handler.go +++ b/go/apps/api/routes/v2_identities_get_identity/handler.go @@ -51,28 +51,26 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - results, err := db.Query.FindIdentityWithRatelimits(ctx, h.DB.RO(), db.FindIdentityWithRatelimitsParams{ + identity, err := db.Query.FindIdentity(ctx, h.DB.RO(), db.FindIdentityParams{ WorkspaceID: auth.AuthorizedWorkspaceID, Identity: req.Identity, Deleted: false, }) if err != nil { + if db.IsNotFound(err) { + return fault.New("identity not found", + fault.Code(codes.Data.Identity.NotFound.URN()), + fault.Internal("identity not found"), + fault.Public("This identity does not exist."), + ) + } + return fault.Wrap(err, fault.Internal("unable to find identity"), fault.Public("We're unable to retrieve the identity."), ) } - if len(results) == 0 { - return fault.New("identity not found", - fault.Code(codes.Data.Identity.NotFound.URN()), - fault.Internal("identity not found"), - fault.Public("This identity does not exist."), - ) - } - - identity := results[0] - // Parse ratelimits JSON var ratelimits []db.RatelimitInfo if ratelimitBytes, ok := identity.Ratelimits.([]byte); ok && ratelimitBytes != nil { diff --git a/go/apps/api/routes/v2_identities_update_identity/handler.go b/go/apps/api/routes/v2_identities_update_identity/handler.go index 9f91a89daf..41571f21d7 100644 --- a/go/apps/api/routes/v2_identities_update_identity/handler.go +++ b/go/apps/api/routes/v2_identities_update_identity/handler.go @@ -113,28 +113,26 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } // Use UNION query to find identity + ratelimits in one query (fast!) - results, err := db.Query.FindIdentityWithRatelimits(ctx, h.DB.RO(), db.FindIdentityWithRatelimitsParams{ + identityRow, err := db.Query.FindIdentity(ctx, h.DB.RO(), db.FindIdentityParams{ WorkspaceID: auth.AuthorizedWorkspaceID, Identity: req.Identity, Deleted: false, }) if err != nil { + if db.IsNotFound(err) { + return fault.New("identity not found", + fault.Code(codes.Data.Identity.NotFound.URN()), + fault.Internal("identity not found"), + fault.Public("Identity not found in this workspace"), + ) + } + return fault.Wrap(err, fault.Internal("unable to find identity"), fault.Public("We're unable to retrieve the identity."), ) } - if len(results) == 0 { - return fault.New("identity not found", - fault.Code(codes.Data.Identity.NotFound.URN()), - fault.Internal("identity not found"), - fault.Public("Identity not found in this workspace"), - ) - } - - identityRow := results[0] - // Parse existing ratelimits from JSON var existingRatelimits []db.RatelimitInfo if ratelimitBytes, ok := identityRow.Ratelimits.([]byte); ok && ratelimitBytes != nil { @@ -142,7 +140,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } type txResult struct { - identity db.FindIdentityWithRatelimitsRow + identity db.FindIdentityRow finalRatelimits []openapi.RatelimitResponse } diff --git a/go/internal/services/keys/verifier.go b/go/internal/services/keys/verifier.go index 242b09f40e..d3453c7629 100644 --- a/go/internal/services/keys/verifier.go +++ b/go/internal/services/keys/verifier.go @@ -48,6 +48,10 @@ type KeyVerifier struct { session *zen.Session // The current request session region string // Geographic region identifier + // Credits - identity credits take priority over key credits + IdentityCredits *db.Credit // Identity-level credits (shared across all keys for this identity) + KeyCredits *db.Credit // Key-level credits (specific to this key) + // Services rateLimiter ratelimit.Service // Rate limiting service usageLimiter usagelimiter.Service // Usage limiting service diff --git a/go/pkg/db/bulk_credits_insert.sql_generated.go b/go/pkg/db/bulk_credits_insert.sql_generated.go new file mode 100644 index 0000000000..9cbfe138e3 --- /dev/null +++ b/go/pkg/db/bulk_credits_insert.sql_generated.go @@ -0,0 +1,47 @@ +// Code generated by sqlc bulk insert plugin. DO NOT EDIT. + +package db + +import ( + "context" + "fmt" + "strings" +) + +// bulkInsertCredit is the base query for bulk insert +const bulkInsertCredit = `INSERT INTO ` + "`" + `credits` + "`" + ` ( id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, created_at, updated_at, refilled_at ) VALUES %s` + +// InsertCredits performs bulk insert in a single query +func (q *BulkQueries) InsertCredits(ctx context.Context, db DBTX, args []InsertCreditParams) error { + + if len(args) == 0 { + return nil + } + + // Build the bulk insert query + valueClauses := make([]string, len(args)) + for i := range args { + valueClauses[i] = "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + } + + bulkQuery := fmt.Sprintf(bulkInsertCredit, strings.Join(valueClauses, ", ")) + + // Collect all arguments + var allArgs []any + for _, arg := range args { + allArgs = append(allArgs, arg.ID) + allArgs = append(allArgs, arg.WorkspaceID) + allArgs = append(allArgs, arg.KeyID) + allArgs = append(allArgs, arg.IdentityID) + allArgs = append(allArgs, arg.Remaining) + allArgs = append(allArgs, arg.RefillDay) + allArgs = append(allArgs, arg.RefillAmount) + allArgs = append(allArgs, arg.CreatedAt) + allArgs = append(allArgs, arg.UpdatedAt) + allArgs = append(allArgs, arg.RefilledAt) + } + + // Execute the bulk insert + _, err := db.ExecContext(ctx, bulkQuery, allArgs...) + return err +} diff --git a/go/pkg/db/bulk_credits_upsert.sql_generated.go b/go/pkg/db/bulk_credits_upsert.sql_generated.go new file mode 100644 index 0000000000..a3e7804310 --- /dev/null +++ b/go/pkg/db/bulk_credits_upsert.sql_generated.go @@ -0,0 +1,72 @@ +// Code generated by sqlc bulk insert plugin. DO NOT EDIT. + +package db + +import ( + "context" + "fmt" + "strings" +) + +// bulkUpsertCredit is the base query for bulk insert +const bulkUpsertCredit = `INSERT INTO ` + "`" + `credits` + "`" + ` ( id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, created_at, updated_at, refilled_at ) VALUES %s ON DUPLICATE KEY UPDATE + remaining = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(remaining) + ELSE remaining + END, + refill_day = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_day) + ELSE refill_day + END, + refill_amount = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_amount) + ELSE refill_amount + END, + refilled_at = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refilled_at) + ELSE refilled_at + END, + updated_at = VALUES(updated_at)` + +// UpsertCredit performs bulk insert in a single query +func (q *BulkQueries) UpsertCredit(ctx context.Context, db DBTX, args []UpsertCreditParams) error { + + if len(args) == 0 { + return nil + } + + // Build the bulk insert query + valueClauses := make([]string, len(args)) + for i := range args { + valueClauses[i] = "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" + } + + bulkQuery := fmt.Sprintf(bulkUpsertCredit, strings.Join(valueClauses, ", ")) + + // Collect all arguments + var allArgs []any + for _, arg := range args { + allArgs = append(allArgs, arg.ID) + allArgs = append(allArgs, arg.WorkspaceID) + allArgs = append(allArgs, arg.KeyID) + allArgs = append(allArgs, arg.IdentityID) + allArgs = append(allArgs, arg.Remaining) + allArgs = append(allArgs, arg.RefillDay) + allArgs = append(allArgs, arg.RefillAmount) + allArgs = append(allArgs, arg.CreatedAt) + allArgs = append(allArgs, arg.UpdatedAt) + allArgs = append(allArgs, arg.RefilledAt) + } + + // Add ON DUPLICATE KEY UPDATE parameters (only once, not per row) + if len(args) > 0 { + allArgs = append(allArgs, args[0].RemainingSpecified) + allArgs = append(allArgs, args[0].RefillDaySpecified) + allArgs = append(allArgs, args[0].RefillAmountSpecified) + allArgs = append(allArgs, args[0].RefilledAtSpecified) + } + + // Execute the bulk insert + _, err := db.ExecContext(ctx, bulkQuery, allArgs...) + return err +} diff --git a/go/pkg/db/credits_delete.sql_generated.go b/go/pkg/db/credits_delete.sql_generated.go new file mode 100644 index 0000000000..a272661a29 --- /dev/null +++ b/go/pkg/db/credits_delete.sql_generated.go @@ -0,0 +1,22 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_delete.sql + +package db + +import ( + "context" +) + +const deleteCredit = `-- name: DeleteCredit :exec +DELETE FROM credits WHERE id = ? +` + +// DeleteCredit +// +// DELETE FROM credits WHERE id = ? +func (q *Queries) DeleteCredit(ctx context.Context, db DBTX, id string) error { + _, err := db.ExecContext(ctx, deleteCredit, id) + return err +} diff --git a/go/pkg/db/credits_find_by_identity_id.sql_generated.go b/go/pkg/db/credits_find_by_identity_id.sql_generated.go new file mode 100644 index 0000000000..cb055cf7f3 --- /dev/null +++ b/go/pkg/db/credits_find_by_identity_id.sql_generated.go @@ -0,0 +1,36 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_find_by_identity_id.sql + +package db + +import ( + "context" + "database/sql" +) + +const findCreditsByIdentityID = `-- name: FindCreditsByIdentityID :one +SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM ` + "`" + `credits` + "`" + ` WHERE identity_id = ? +` + +// FindCreditsByIdentityID +// +// SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM `credits` WHERE identity_id = ? +func (q *Queries) FindCreditsByIdentityID(ctx context.Context, db DBTX, identityID sql.NullString) (Credit, error) { + row := db.QueryRowContext(ctx, findCreditsByIdentityID, identityID) + var i Credit + err := row.Scan( + &i.ID, + &i.WorkspaceID, + &i.KeyID, + &i.IdentityID, + &i.Remaining, + &i.RefillDay, + &i.RefillAmount, + &i.RefilledAt, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} diff --git a/go/pkg/db/credits_find_by_key_id.sql_generated.go b/go/pkg/db/credits_find_by_key_id.sql_generated.go new file mode 100644 index 0000000000..7cf1a17518 --- /dev/null +++ b/go/pkg/db/credits_find_by_key_id.sql_generated.go @@ -0,0 +1,36 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_find_by_key_id.sql + +package db + +import ( + "context" + "database/sql" +) + +const findCreditsByKeyID = `-- name: FindCreditsByKeyID :one +SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM ` + "`" + `credits` + "`" + ` WHERE key_id = ? +` + +// FindCreditsByKeyID +// +// SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM `credits` WHERE key_id = ? +func (q *Queries) FindCreditsByKeyID(ctx context.Context, db DBTX, keyID sql.NullString) (Credit, error) { + row := db.QueryRowContext(ctx, findCreditsByKeyID, keyID) + var i Credit + err := row.Scan( + &i.ID, + &i.WorkspaceID, + &i.KeyID, + &i.IdentityID, + &i.Remaining, + &i.RefillDay, + &i.RefillAmount, + &i.RefilledAt, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} diff --git a/go/pkg/db/credits_find_remaining.sql_generated.go b/go/pkg/db/credits_find_remaining.sql_generated.go new file mode 100644 index 0000000000..e4e5eaa487 --- /dev/null +++ b/go/pkg/db/credits_find_remaining.sql_generated.go @@ -0,0 +1,24 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_find_remaining.sql + +package db + +import ( + "context" +) + +const findRemainingCredits = `-- name: FindRemainingCredits :one +SELECT remaining FROM ` + "`" + `credits` + "`" + ` WHERE id = ? +` + +// FindRemainingCredits +// +// SELECT remaining FROM `credits` WHERE id = ? +func (q *Queries) FindRemainingCredits(ctx context.Context, db DBTX, id string) (int32, error) { + row := db.QueryRowContext(ctx, findRemainingCredits, id) + var remaining int32 + err := row.Scan(&remaining) + return remaining, err +} diff --git a/go/pkg/db/credits_insert.sql_generated.go b/go/pkg/db/credits_insert.sql_generated.go new file mode 100644 index 0000000000..a07e28098c --- /dev/null +++ b/go/pkg/db/credits_insert.sql_generated.go @@ -0,0 +1,69 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_insert.sql + +package db + +import ( + "context" + "database/sql" +) + +const insertCredit = `-- name: InsertCredit :exec +INSERT INTO ` + "`" + `credits` + "`" + ` ( + id, + workspace_id, + key_id, + identity_id, + remaining, + refill_day, + refill_amount, + created_at, + updated_at, + refilled_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +` + +type InsertCreditParams struct { + ID string `db:"id"` + WorkspaceID string `db:"workspace_id"` + KeyID sql.NullString `db:"key_id"` + IdentityID sql.NullString `db:"identity_id"` + Remaining int32 `db:"remaining"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + CreatedAt int64 `db:"created_at"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + RefilledAt sql.NullInt64 `db:"refilled_at"` +} + +// InsertCredit +// +// INSERT INTO `credits` ( +// id, +// workspace_id, +// key_id, +// identity_id, +// remaining, +// refill_day, +// refill_amount, +// created_at, +// updated_at, +// refilled_at +// ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +func (q *Queries) InsertCredit(ctx context.Context, db DBTX, arg InsertCreditParams) error { + _, err := db.ExecContext(ctx, insertCredit, + arg.ID, + arg.WorkspaceID, + arg.KeyID, + arg.IdentityID, + arg.Remaining, + arg.RefillDay, + arg.RefillAmount, + arg.CreatedAt, + arg.UpdatedAt, + arg.RefilledAt, + ) + return err +} diff --git a/go/pkg/db/credits_update_decrement.sql_generated.go b/go/pkg/db/credits_update_decrement.sql_generated.go new file mode 100644 index 0000000000..5f4fa3b8be --- /dev/null +++ b/go/pkg/db/credits_update_decrement.sql_generated.go @@ -0,0 +1,46 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_update_decrement.sql + +package db + +import ( + "context" + "database/sql" +) + +const updateCreditDecrement = `-- name: UpdateCreditDecrement :exec +UPDATE ` + "`" + `credits` + "`" + ` +SET remaining = CASE + WHEN remaining >= ? THEN remaining - ? + ELSE 0 +END, + updated_at = ? +WHERE id = ? +` + +type UpdateCreditDecrementParams struct { + Credits int32 `db:"credits"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + ID string `db:"id"` +} + +// UpdateCreditDecrement +// +// UPDATE `credits` +// SET remaining = CASE +// WHEN remaining >= ? THEN remaining - ? +// ELSE 0 +// END, +// updated_at = ? +// WHERE id = ? +func (q *Queries) UpdateCreditDecrement(ctx context.Context, db DBTX, arg UpdateCreditDecrementParams) error { + _, err := db.ExecContext(ctx, updateCreditDecrement, + arg.Credits, + arg.Credits, + arg.UpdatedAt, + arg.ID, + ) + return err +} diff --git a/go/pkg/db/credits_update_increment.sql_generated.go b/go/pkg/db/credits_update_increment.sql_generated.go new file mode 100644 index 0000000000..3e9e7e89f2 --- /dev/null +++ b/go/pkg/db/credits_update_increment.sql_generated.go @@ -0,0 +1,35 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_update_increment.sql + +package db + +import ( + "context" + "database/sql" +) + +const updateCreditIncrement = `-- name: UpdateCreditIncrement :exec +UPDATE ` + "`" + `credits` + "`" + ` +SET remaining = remaining + ?, + updated_at = ? +WHERE id = ? +` + +type UpdateCreditIncrementParams struct { + Credits int32 `db:"credits"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + ID string `db:"id"` +} + +// UpdateCreditIncrement +// +// UPDATE `credits` +// SET remaining = remaining + ?, +// updated_at = ? +// WHERE id = ? +func (q *Queries) UpdateCreditIncrement(ctx context.Context, db DBTX, arg UpdateCreditIncrementParams) error { + _, err := db.ExecContext(ctx, updateCreditIncrement, arg.Credits, arg.UpdatedAt, arg.ID) + return err +} diff --git a/go/pkg/db/credits_update_set.sql_generated.go b/go/pkg/db/credits_update_set.sql_generated.go new file mode 100644 index 0000000000..8a7056a0e3 --- /dev/null +++ b/go/pkg/db/credits_update_set.sql_generated.go @@ -0,0 +1,35 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_update_set.sql + +package db + +import ( + "context" + "database/sql" +) + +const updateCreditSet = `-- name: UpdateCreditSet :exec +UPDATE ` + "`" + `credits` + "`" + ` +SET remaining = ?, + updated_at = ? +WHERE id = ? +` + +type UpdateCreditSetParams struct { + Remaining int32 `db:"remaining"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + ID string `db:"id"` +} + +// UpdateCreditSet +// +// UPDATE `credits` +// SET remaining = ?, +// updated_at = ? +// WHERE id = ? +func (q *Queries) UpdateCreditSet(ctx context.Context, db DBTX, arg UpdateCreditSetParams) error { + _, err := db.ExecContext(ctx, updateCreditSet, arg.Remaining, arg.UpdatedAt, arg.ID) + return err +} diff --git a/go/pkg/db/credits_upsert.sql_generated.go b/go/pkg/db/credits_upsert.sql_generated.go new file mode 100644 index 0000000000..5581f39a4d --- /dev/null +++ b/go/pkg/db/credits_upsert.sql_generated.go @@ -0,0 +1,113 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: credits_upsert.sql + +package db + +import ( + "context" + "database/sql" +) + +const upsertCredit = `-- name: UpsertCredit :exec +INSERT INTO ` + "`" + `credits` + "`" + ` ( + id, + workspace_id, + key_id, + identity_id, + remaining, + refill_day, + refill_amount, + created_at, + updated_at, + refilled_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + remaining = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(remaining) + ELSE remaining + END, + refill_day = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_day) + ELSE refill_day + END, + refill_amount = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_amount) + ELSE refill_amount + END, + refilled_at = CASE + WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refilled_at) + ELSE refilled_at + END, + updated_at = VALUES(updated_at) +` + +type UpsertCreditParams struct { + ID string `db:"id"` + WorkspaceID string `db:"workspace_id"` + KeyID sql.NullString `db:"key_id"` + IdentityID sql.NullString `db:"identity_id"` + Remaining int32 `db:"remaining"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + CreatedAt int64 `db:"created_at"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + RefilledAt sql.NullInt64 `db:"refilled_at"` + RemainingSpecified int64 `db:"remaining_specified"` + RefillDaySpecified int64 `db:"refill_day_specified"` + RefillAmountSpecified int64 `db:"refill_amount_specified"` + RefilledAtSpecified int64 `db:"refilled_at_specified"` +} + +// UpsertCredit +// +// INSERT INTO `credits` ( +// id, +// workspace_id, +// key_id, +// identity_id, +// remaining, +// refill_day, +// refill_amount, +// created_at, +// updated_at, +// refilled_at +// ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +// ON DUPLICATE KEY UPDATE +// remaining = CASE +// WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(remaining) +// ELSE remaining +// END, +// refill_day = CASE +// WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_day) +// ELSE refill_day +// END, +// refill_amount = CASE +// WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_amount) +// ELSE refill_amount +// END, +// refilled_at = CASE +// WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refilled_at) +// ELSE refilled_at +// END, +// updated_at = VALUES(updated_at) +func (q *Queries) UpsertCredit(ctx context.Context, db DBTX, arg UpsertCreditParams) error { + _, err := db.ExecContext(ctx, upsertCredit, + arg.ID, + arg.WorkspaceID, + arg.KeyID, + arg.IdentityID, + arg.Remaining, + arg.RefillDay, + arg.RefillAmount, + arg.CreatedAt, + arg.UpdatedAt, + arg.RefilledAt, + arg.RemainingSpecified, + arg.RefillDaySpecified, + arg.RefillAmountSpecified, + arg.RefilledAtSpecified, + ) + return err +} diff --git a/go/pkg/db/identity_delete_old_by_external_id.sql_generated.go b/go/pkg/db/identity_delete_old_by_external_id.sql_generated.go index cf0535d3c4..61e08fb148 100644 --- a/go/pkg/db/identity_delete_old_by_external_id.sql_generated.go +++ b/go/pkg/db/identity_delete_old_by_external_id.sql_generated.go @@ -10,9 +10,10 @@ import ( ) const deleteOldIdentityByExternalID = `-- name: DeleteOldIdentityByExternalID :exec -DELETE i, rl +DELETE i, rl, c FROM identities i LEFT JOIN ratelimits rl ON rl.identity_id = i.id +LEFT JOIN credits c ON c.identity_id = i.id WHERE i.workspace_id = ? AND i.external_id = ? AND i.id != ? @@ -27,9 +28,10 @@ type DeleteOldIdentityByExternalIDParams struct { // DeleteOldIdentityByExternalID // -// DELETE i, rl +// DELETE i, rl, c // FROM identities i // LEFT JOIN ratelimits rl ON rl.identity_id = i.id +// LEFT JOIN credits c ON c.identity_id = i.id // WHERE i.workspace_id = ? // AND i.external_id = ? // AND i.id != ? diff --git a/go/pkg/db/identity_delete_old_with_ratelimits.sql_generated.go b/go/pkg/db/identity_delete_old_with_ratelimits.sql_generated.go index 344bde2a3c..c57f973535 100644 --- a/go/pkg/db/identity_delete_old_with_ratelimits.sql_generated.go +++ b/go/pkg/db/identity_delete_old_with_ratelimits.sql_generated.go @@ -10,9 +10,10 @@ import ( ) const deleteOldIdentityWithRatelimits = `-- name: DeleteOldIdentityWithRatelimits :exec -DELETE i, rl +DELETE i, rl, c FROM identities i LEFT JOIN ratelimits rl ON rl.identity_id = i.id +LEFT JOIN credits c ON c.identity_id = i.id WHERE i.workspace_id = ? AND (i.id = ? OR i.external_id = ?) AND i.deleted = true @@ -25,9 +26,10 @@ type DeleteOldIdentityWithRatelimitsParams struct { // DeleteOldIdentityWithRatelimits // -// DELETE i, rl +// DELETE i, rl, c // FROM identities i // LEFT JOIN ratelimits rl ON rl.identity_id = i.id +// LEFT JOIN credits c ON c.identity_id = i.id // WHERE i.workspace_id = ? // AND (i.id = ? OR i.external_id = ?) // AND i.deleted = true diff --git a/go/pkg/db/identity_find.sql_generated.go b/go/pkg/db/identity_find.sql_generated.go new file mode 100644 index 0000000000..e28f4d20cd --- /dev/null +++ b/go/pkg/db/identity_find.sql_generated.go @@ -0,0 +1,140 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: identity_find.sql + +package db + +import ( + "context" + "database/sql" +) + +const findIdentity = `-- name: FindIdentity :one +SELECT + i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, + c.id AS credit_id, + c.remaining AS credit_remaining, + c.refill_amount AS credit_refill_amount, + c.refill_day AS credit_refill_day, + c.refilled_at AS credit_refilled_at, + COALESCE( + (SELECT JSON_ARRAYAGG( + JSON_OBJECT( + 'id', rl.id, + 'name', rl.name, + 'key_id', rl.key_id, + 'identity_id', rl.identity_id, + 'limit', rl.` + "`" + `limit` + "`" + `, + 'duration', rl.duration, + 'auto_apply', rl.auto_apply = 1 + ) + ) + FROM ratelimits rl WHERE rl.identity_id = i.id), + JSON_ARRAY() + ) as ratelimits +FROM identities i +LEFT JOIN credits c ON c.identity_id = i.id +WHERE i.id = ( + SELECT id FROM identities sub1 + WHERE sub1.workspace_id = ? + AND sub1.id = ? + AND sub1.deleted = ? + UNION ALL + SELECT id FROM identities sub2 + WHERE sub2.workspace_id = ? + AND sub2.external_id = ? + AND sub2.deleted = ? + LIMIT 1 +) +` + +type FindIdentityParams struct { + WorkspaceID string `db:"workspace_id"` + Identity string `db:"identity"` + Deleted bool `db:"deleted"` +} + +type FindIdentityRow struct { + ID string `db:"id"` + ExternalID string `db:"external_id"` + WorkspaceID string `db:"workspace_id"` + Environment string `db:"environment"` + Meta []byte `db:"meta"` + Deleted bool `db:"deleted"` + CreatedAt int64 `db:"created_at"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + CreditID sql.NullString `db:"credit_id"` + CreditRemaining sql.NullInt32 `db:"credit_remaining"` + CreditRefillAmount sql.NullInt32 `db:"credit_refill_amount"` + CreditRefillDay sql.NullInt16 `db:"credit_refill_day"` + CreditRefilledAt sql.NullInt64 `db:"credit_refilled_at"` + Ratelimits interface{} `db:"ratelimits"` +} + +// FindIdentity +// +// SELECT +// i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, +// c.id AS credit_id, +// c.remaining AS credit_remaining, +// c.refill_amount AS credit_refill_amount, +// c.refill_day AS credit_refill_day, +// c.refilled_at AS credit_refilled_at, +// COALESCE( +// (SELECT JSON_ARRAYAGG( +// JSON_OBJECT( +// 'id', rl.id, +// 'name', rl.name, +// 'key_id', rl.key_id, +// 'identity_id', rl.identity_id, +// 'limit', rl.`limit`, +// 'duration', rl.duration, +// 'auto_apply', rl.auto_apply = 1 +// ) +// ) +// FROM ratelimits rl WHERE rl.identity_id = i.id), +// JSON_ARRAY() +// ) as ratelimits +// FROM identities i +// LEFT JOIN credits c ON c.identity_id = i.id +// WHERE i.id = ( +// SELECT id FROM identities sub1 +// WHERE sub1.workspace_id = ? +// AND sub1.id = ? +// AND sub1.deleted = ? +// UNION ALL +// SELECT id FROM identities sub2 +// WHERE sub2.workspace_id = ? +// AND sub2.external_id = ? +// AND sub2.deleted = ? +// LIMIT 1 +// ) +func (q *Queries) FindIdentity(ctx context.Context, db DBTX, arg FindIdentityParams) (FindIdentityRow, error) { + row := db.QueryRowContext(ctx, findIdentity, + arg.WorkspaceID, + arg.Identity, + arg.Deleted, + arg.WorkspaceID, + arg.Identity, + arg.Deleted, + ) + var i FindIdentityRow + err := row.Scan( + &i.ID, + &i.ExternalID, + &i.WorkspaceID, + &i.Environment, + &i.Meta, + &i.Deleted, + &i.CreatedAt, + &i.UpdatedAt, + &i.CreditID, + &i.CreditRemaining, + &i.CreditRefillAmount, + &i.CreditRefillDay, + &i.CreditRefilledAt, + &i.Ratelimits, + ) + return i, err +} diff --git a/go/pkg/db/identity_find_with_ratelimits.sql_generated.go b/go/pkg/db/identity_find_with_ratelimits.sql_generated.go deleted file mode 100644 index 92558675b4..0000000000 --- a/go/pkg/db/identity_find_with_ratelimits.sql_generated.go +++ /dev/null @@ -1,162 +0,0 @@ -// Code generated by sqlc. DO NOT EDIT. -// versions: -// sqlc v1.29.0 -// source: identity_find_with_ratelimits.sql - -package db - -import ( - "context" - "database/sql" -) - -const findIdentityWithRatelimits = `-- name: FindIdentityWithRatelimits :many -SELECT - i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, - COALESCE( - (SELECT JSON_ARRAYAGG( - JSON_OBJECT( - 'id', rl.id, - 'name', rl.name, - 'key_id', rl.key_id, - 'identity_id', rl.identity_id, - 'limit', rl.` + "`" + `limit` + "`" + `, - 'duration', rl.duration, - 'auto_apply', rl.auto_apply = 1 - ) - ) - FROM ratelimits rl WHERE rl.identity_id = i.id), - JSON_ARRAY() - ) as ratelimits -FROM identities i -WHERE i.workspace_id = ? - AND i.id = ? - AND i.deleted = ? -UNION ALL -SELECT - i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, - COALESCE( - (SELECT JSON_ARRAYAGG( - JSON_OBJECT( - 'id', rl.id, - 'name', rl.name, - 'key_id', rl.key_id, - 'identity_id', rl.identity_id, - 'limit', rl.` + "`" + `limit` + "`" + `, - 'duration', rl.duration, - 'auto_apply', rl.auto_apply = 1 - ) - ) - FROM ratelimits rl WHERE rl.identity_id = i.id), - JSON_ARRAY() - ) as ratelimits -FROM identities i -WHERE i.workspace_id = ? - AND i.external_id = ? - AND i.deleted = ? -LIMIT 1 -` - -type FindIdentityWithRatelimitsParams struct { - WorkspaceID string `db:"workspace_id"` - Identity string `db:"identity"` - Deleted bool `db:"deleted"` -} - -type FindIdentityWithRatelimitsRow struct { - ID string `db:"id"` - ExternalID string `db:"external_id"` - WorkspaceID string `db:"workspace_id"` - Environment string `db:"environment"` - Meta []byte `db:"meta"` - Deleted bool `db:"deleted"` - CreatedAt int64 `db:"created_at"` - UpdatedAt sql.NullInt64 `db:"updated_at"` - Ratelimits interface{} `db:"ratelimits"` -} - -// FindIdentityWithRatelimits -// -// SELECT -// i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, -// COALESCE( -// (SELECT JSON_ARRAYAGG( -// JSON_OBJECT( -// 'id', rl.id, -// 'name', rl.name, -// 'key_id', rl.key_id, -// 'identity_id', rl.identity_id, -// 'limit', rl.`limit`, -// 'duration', rl.duration, -// 'auto_apply', rl.auto_apply = 1 -// ) -// ) -// FROM ratelimits rl WHERE rl.identity_id = i.id), -// JSON_ARRAY() -// ) as ratelimits -// FROM identities i -// WHERE i.workspace_id = ? -// AND i.id = ? -// AND i.deleted = ? -// UNION ALL -// SELECT -// i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, -// COALESCE( -// (SELECT JSON_ARRAYAGG( -// JSON_OBJECT( -// 'id', rl.id, -// 'name', rl.name, -// 'key_id', rl.key_id, -// 'identity_id', rl.identity_id, -// 'limit', rl.`limit`, -// 'duration', rl.duration, -// 'auto_apply', rl.auto_apply = 1 -// ) -// ) -// FROM ratelimits rl WHERE rl.identity_id = i.id), -// JSON_ARRAY() -// ) as ratelimits -// FROM identities i -// WHERE i.workspace_id = ? -// AND i.external_id = ? -// AND i.deleted = ? -// LIMIT 1 -func (q *Queries) FindIdentityWithRatelimits(ctx context.Context, db DBTX, arg FindIdentityWithRatelimitsParams) ([]FindIdentityWithRatelimitsRow, error) { - rows, err := db.QueryContext(ctx, findIdentityWithRatelimits, - arg.WorkspaceID, - arg.Identity, - arg.Deleted, - arg.WorkspaceID, - arg.Identity, - arg.Deleted, - ) - if err != nil { - return nil, err - } - defer rows.Close() - var items []FindIdentityWithRatelimitsRow - for rows.Next() { - var i FindIdentityWithRatelimitsRow - if err := rows.Scan( - &i.ID, - &i.ExternalID, - &i.WorkspaceID, - &i.Environment, - &i.Meta, - &i.Deleted, - &i.CreatedAt, - &i.UpdatedAt, - &i.Ratelimits, - ); err != nil { - return nil, err - } - items = append(items, i) - } - if err := rows.Close(); err != nil { - return nil, err - } - if err := rows.Err(); err != nil { - return nil, err - } - return items, nil -} diff --git a/go/pkg/db/identity_list.sql_generated.go b/go/pkg/db/identity_list.sql_generated.go index 1e0643dd20..cb56485f2d 100644 --- a/go/pkg/db/identity_list.sql_generated.go +++ b/go/pkg/db/identity_list.sql_generated.go @@ -7,15 +7,22 @@ package db import ( "context" + "database/sql" ) const listIdentities = `-- name: ListIdentities :many -SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at -FROM identities -WHERE workspace_id = ? -AND deleted = ? -AND id >= ? -ORDER BY id ASC +SELECT + i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, + c.id as credit_id, + c.remaining as credit_remaining, + c.refill_amount as credit_refill_amount, + c.refill_day as credit_refill_day +FROM identities i +LEFT JOIN credits c ON c.identity_id = i.id +WHERE i.workspace_id = ? +AND i.deleted = ? +AND i.id >= ? +ORDER BY i.id ASC LIMIT ? ` @@ -26,16 +33,37 @@ type ListIdentitiesParams struct { Limit int32 `db:"limit"` } +type ListIdentitiesRow struct { + ID string `db:"id"` + ExternalID string `db:"external_id"` + WorkspaceID string `db:"workspace_id"` + Environment string `db:"environment"` + Meta []byte `db:"meta"` + Deleted bool `db:"deleted"` + CreatedAt int64 `db:"created_at"` + UpdatedAt sql.NullInt64 `db:"updated_at"` + CreditID sql.NullString `db:"credit_id"` + CreditRemaining sql.NullInt32 `db:"credit_remaining"` + CreditRefillAmount sql.NullInt32 `db:"credit_refill_amount"` + CreditRefillDay sql.NullInt16 `db:"credit_refill_day"` +} + // ListIdentities // -// SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at -// FROM identities -// WHERE workspace_id = ? -// AND deleted = ? -// AND id >= ? -// ORDER BY id ASC +// SELECT +// i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, +// c.id as credit_id, +// c.remaining as credit_remaining, +// c.refill_amount as credit_refill_amount, +// c.refill_day as credit_refill_day +// FROM identities i +// LEFT JOIN credits c ON c.identity_id = i.id +// WHERE i.workspace_id = ? +// AND i.deleted = ? +// AND i.id >= ? +// ORDER BY i.id ASC // LIMIT ? -func (q *Queries) ListIdentities(ctx context.Context, db DBTX, arg ListIdentitiesParams) ([]Identity, error) { +func (q *Queries) ListIdentities(ctx context.Context, db DBTX, arg ListIdentitiesParams) ([]ListIdentitiesRow, error) { rows, err := db.QueryContext(ctx, listIdentities, arg.WorkspaceID, arg.Deleted, @@ -46,9 +74,9 @@ func (q *Queries) ListIdentities(ctx context.Context, db DBTX, arg ListIdentitie return nil, err } defer rows.Close() - var items []Identity + var items []ListIdentitiesRow for rows.Next() { - var i Identity + var i ListIdentitiesRow if err := rows.Scan( &i.ID, &i.ExternalID, @@ -58,6 +86,10 @@ func (q *Queries) ListIdentities(ctx context.Context, db DBTX, arg ListIdentitie &i.Deleted, &i.CreatedAt, &i.UpdatedAt, + &i.CreditID, + &i.CreditRemaining, + &i.CreditRefillAmount, + &i.CreditRefillDay, ); err != nil { return nil, err } diff --git a/go/pkg/db/key_data.go b/go/pkg/db/key_data.go index 38a22dd181..bc3cc75854 100644 --- a/go/pkg/db/key_data.go +++ b/go/pkg/db/key_data.go @@ -12,6 +12,8 @@ type KeyData struct { KeyAuth KeyAuth Workspace Workspace Identity *Identity // Is optional + KeyCredits *Credit // Credits associated with the key + IdentityCredits *Credit // Credits associated with the identity EncryptedKey sql.NullString EncryptionKeyID sql.NullString Roles []RoleInfo @@ -170,6 +172,34 @@ func buildKeyData(r *FindLiveKeyByHashRow) *KeyData { } } + // Populate key credits if they exist + if r.CreditID.Valid { + kd.KeyCredits = &Credit{ + ID: r.CreditID.String, + WorkspaceID: r.WorkspaceID, + KeyID: sql.NullString{Valid: true, String: r.ID}, + IdentityID: sql.NullString{Valid: false}, + Remaining: r.CreditRemaining.Int32, + RefillDay: r.CreditRefillDay, + RefillAmount: r.CreditRefillAmount, + RefilledAt: r.CreditRefilledAt, + } + } + + // Populate identity credits if they exist + if r.IdentityCreditID.Valid { + kd.IdentityCredits = &Credit{ + ID: r.IdentityCreditID.String, + WorkspaceID: r.WorkspaceID, + KeyID: sql.NullString{Valid: false}, + IdentityID: r.IdentityID, + Remaining: r.IdentityCreditRemaining.Int32, + RefillDay: r.IdentityCreditRefillDay, + RefillAmount: r.IdentityCreditRefillAmount, + RefilledAt: r.IdentityCreditRefilledAt, + } + } + // It's fine to fail here if roleBytes, ok := r.Roles.([]byte); ok && roleBytes != nil { _ = json.Unmarshal(roleBytes, &kd.Roles) // Ignore error, default to empty array diff --git a/go/pkg/db/key_find_for_verification.sql_generated.go b/go/pkg/db/key_find_for_verification.sql_generated.go index 1093c89159..b3b9928ced 100644 --- a/go/pkg/db/key_find_for_verification.sql_generated.go +++ b/go/pkg/db/key_find_for_verification.sql_generated.go @@ -78,44 +78,71 @@ select k.id, i.meta as identity_meta, ka.deleted_at_m as key_auth_deleted_at_m, ws.enabled as workspace_enabled, - fws.enabled as for_workspace_enabled + fws.enabled as for_workspace_enabled, + + -- Key-level credits + key_credits.id as key_credit_id, + key_credits.remaining as key_credit_remaining, + key_credits.refill_amount as key_credit_refill_amount, + key_credits.refill_day as key_credit_refill_day, + key_credits.refilled_at as key_credit_refilled_at, + + -- Identity-level credits + identity_credits.id as identity_credit_id, + identity_credits.remaining as identity_credit_remaining, + identity_credits.refill_amount as identity_credit_refill_amount, + identity_credits.refill_day as identity_credit_refill_day, + identity_credits.refilled_at as identity_credit_refilled_at + from ` + "`" + `keys` + "`" + ` k JOIN apis a USING (key_auth_id) JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN workspaces fws ON fws.id = k.for_workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = 0 + LEFT JOIN credits key_credits ON key_credits.key_id = k.id + LEFT JOIN credits identity_credits ON identity_credits.identity_id = i.id where k.hash = ? and k.deleted_at_m is null ` type FindKeyForVerificationRow struct { - ID string `db:"id"` - KeyAuthID string `db:"key_auth_id"` - WorkspaceID string `db:"workspace_id"` - ForWorkspaceID sql.NullString `db:"for_workspace_id"` - Name sql.NullString `db:"name"` - Meta sql.NullString `db:"meta"` - Expires sql.NullTime `db:"expires"` - DeletedAtM sql.NullInt64 `db:"deleted_at_m"` - RefillDay sql.NullInt16 `db:"refill_day"` - RefillAmount sql.NullInt32 `db:"refill_amount"` - LastRefillAt sql.NullTime `db:"last_refill_at"` - Enabled bool `db:"enabled"` - RemainingRequests sql.NullInt32 `db:"remaining_requests"` - IpWhitelist sql.NullString `db:"ip_whitelist"` - ApiWorkspaceID string `db:"api_workspace_id"` - ApiID string `db:"api_id"` - ApiDeletedAtM sql.NullInt64 `db:"api_deleted_at_m"` - Roles interface{} `db:"roles"` - Permissions interface{} `db:"permissions"` - Ratelimits interface{} `db:"ratelimits"` - IdentityID sql.NullString `db:"identity_id"` - ExternalID sql.NullString `db:"external_id"` - IdentityMeta []byte `db:"identity_meta"` - KeyAuthDeletedAtM sql.NullInt64 `db:"key_auth_deleted_at_m"` - WorkspaceEnabled bool `db:"workspace_enabled"` - ForWorkspaceEnabled sql.NullBool `db:"for_workspace_enabled"` + ID string `db:"id"` + KeyAuthID string `db:"key_auth_id"` + WorkspaceID string `db:"workspace_id"` + ForWorkspaceID sql.NullString `db:"for_workspace_id"` + Name sql.NullString `db:"name"` + Meta sql.NullString `db:"meta"` + Expires sql.NullTime `db:"expires"` + DeletedAtM sql.NullInt64 `db:"deleted_at_m"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + LastRefillAt sql.NullTime `db:"last_refill_at"` + Enabled bool `db:"enabled"` + RemainingRequests sql.NullInt32 `db:"remaining_requests"` + IpWhitelist sql.NullString `db:"ip_whitelist"` + ApiWorkspaceID string `db:"api_workspace_id"` + ApiID string `db:"api_id"` + ApiDeletedAtM sql.NullInt64 `db:"api_deleted_at_m"` + Roles interface{} `db:"roles"` + Permissions interface{} `db:"permissions"` + Ratelimits interface{} `db:"ratelimits"` + IdentityID sql.NullString `db:"identity_id"` + ExternalID sql.NullString `db:"external_id"` + IdentityMeta []byte `db:"identity_meta"` + KeyAuthDeletedAtM sql.NullInt64 `db:"key_auth_deleted_at_m"` + WorkspaceEnabled bool `db:"workspace_enabled"` + ForWorkspaceEnabled sql.NullBool `db:"for_workspace_enabled"` + KeyCreditID sql.NullString `db:"key_credit_id"` + KeyCreditRemaining sql.NullInt32 `db:"key_credit_remaining"` + KeyCreditRefillAmount sql.NullInt32 `db:"key_credit_refill_amount"` + KeyCreditRefillDay sql.NullInt16 `db:"key_credit_refill_day"` + KeyCreditRefilledAt sql.NullInt64 `db:"key_credit_refilled_at"` + IdentityCreditID sql.NullString `db:"identity_credit_id"` + IdentityCreditRemaining sql.NullInt32 `db:"identity_credit_remaining"` + IdentityCreditRefillAmount sql.NullInt32 `db:"identity_credit_refill_amount"` + IdentityCreditRefillDay sql.NullInt16 `db:"identity_credit_refill_day"` + IdentityCreditRefilledAt sql.NullInt64 `db:"identity_credit_refilled_at"` } // FindKeyForVerification @@ -187,13 +214,30 @@ type FindKeyForVerificationRow struct { // i.meta as identity_meta, // ka.deleted_at_m as key_auth_deleted_at_m, // ws.enabled as workspace_enabled, -// fws.enabled as for_workspace_enabled +// fws.enabled as for_workspace_enabled, +// +// -- Key-level credits +// key_credits.id as key_credit_id, +// key_credits.remaining as key_credit_remaining, +// key_credits.refill_amount as key_credit_refill_amount, +// key_credits.refill_day as key_credit_refill_day, +// key_credits.refilled_at as key_credit_refilled_at, +// +// -- Identity-level credits +// identity_credits.id as identity_credit_id, +// identity_credits.remaining as identity_credit_remaining, +// identity_credits.refill_amount as identity_credit_refill_amount, +// identity_credits.refill_day as identity_credit_refill_day, +// identity_credits.refilled_at as identity_credit_refilled_at +// // from `keys` k // JOIN apis a USING (key_auth_id) // JOIN key_auth ka ON ka.id = k.key_auth_id // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN workspaces fws ON fws.id = k.for_workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = 0 +// LEFT JOIN credits key_credits ON key_credits.key_id = k.id +// LEFT JOIN credits identity_credits ON identity_credits.identity_id = i.id // where k.hash = ? // and k.deleted_at_m is null func (q *Queries) FindKeyForVerification(ctx context.Context, db DBTX, hash string) (FindKeyForVerificationRow, error) { @@ -226,6 +270,16 @@ func (q *Queries) FindKeyForVerification(ctx context.Context, db DBTX, hash stri &i.KeyAuthDeletedAtM, &i.WorkspaceEnabled, &i.ForWorkspaceEnabled, + &i.KeyCreditID, + &i.KeyCreditRemaining, + &i.KeyCreditRefillAmount, + &i.KeyCreditRefillDay, + &i.KeyCreditRefilledAt, + &i.IdentityCreditID, + &i.IdentityCreditRemaining, + &i.IdentityCreditRefillAmount, + &i.IdentityCreditRefillDay, + &i.IdentityCreditRefilledAt, ) return i, err } diff --git a/go/pkg/db/key_find_live_by_hash.sql_generated.go b/go/pkg/db/key_find_live_by_hash.sql_generated.go index 1788c94135..36326c4d36 100644 --- a/go/pkg/db/key_find_live_by_hash.sql_generated.go +++ b/go/pkg/db/key_find_live_by_hash.sql_generated.go @@ -86,7 +86,21 @@ SELECT FROM ratelimits rl WHERE rl.key_id = k.id OR rl.identity_id = i.id), JSON_ARRAY() - ) as ratelimits + ) as ratelimits, + + -- Key credits + kc.id as credit_id, + kc.remaining as credit_remaining, + kc.refill_amount as credit_refill_amount, + kc.refill_day as credit_refill_day, + kc.refilled_at as credit_refilled_at, + + -- Identity credits + ic.id as identity_credit_id, + ic.remaining as identity_credit_remaining, + ic.refill_amount as identity_credit_refill_amount, + ic.refill_day as identity_credit_refill_day, + ic.refilled_at as identity_credit_refilled_at FROM ` + "`" + `keys` + "`" + ` k JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -94,6 +108,8 @@ JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +LEFT JOIN credits kc ON kc.key_id = k.id +LEFT JOIN credits ic ON ic.identity_id = i.id WHERE k.hash = ? AND k.deleted_at_m IS NULL AND a.deleted_at_m IS NULL @@ -102,41 +118,51 @@ WHERE k.hash = ? ` type FindLiveKeyByHashRow struct { - ID string `db:"id"` - KeyAuthID string `db:"key_auth_id"` - Hash string `db:"hash"` - Start string `db:"start"` - WorkspaceID string `db:"workspace_id"` - ForWorkspaceID sql.NullString `db:"for_workspace_id"` - Name sql.NullString `db:"name"` - OwnerID sql.NullString `db:"owner_id"` - IdentityID sql.NullString `db:"identity_id"` - Meta sql.NullString `db:"meta"` - Expires sql.NullTime `db:"expires"` - CreatedAtM int64 `db:"created_at_m"` - UpdatedAtM sql.NullInt64 `db:"updated_at_m"` - DeletedAtM sql.NullInt64 `db:"deleted_at_m"` - RefillDay sql.NullInt16 `db:"refill_day"` - RefillAmount sql.NullInt32 `db:"refill_amount"` - LastRefillAt sql.NullTime `db:"last_refill_at"` - Enabled bool `db:"enabled"` - RemainingRequests sql.NullInt32 `db:"remaining_requests"` - RatelimitAsync sql.NullBool `db:"ratelimit_async"` - RatelimitLimit sql.NullInt32 `db:"ratelimit_limit"` - RatelimitDuration sql.NullInt64 `db:"ratelimit_duration"` - Environment sql.NullString `db:"environment"` - Api Api `db:"api"` - KeyAuth KeyAuth `db:"key_auth"` - Workspace Workspace `db:"workspace"` - IdentityTableID sql.NullString `db:"identity_table_id"` - IdentityExternalID sql.NullString `db:"identity_external_id"` - IdentityMeta []byte `db:"identity_meta"` - EncryptedKey sql.NullString `db:"encrypted_key"` - EncryptionKeyID sql.NullString `db:"encryption_key_id"` - Roles interface{} `db:"roles"` - Permissions interface{} `db:"permissions"` - RolePermissions interface{} `db:"role_permissions"` - Ratelimits interface{} `db:"ratelimits"` + ID string `db:"id"` + KeyAuthID string `db:"key_auth_id"` + Hash string `db:"hash"` + Start string `db:"start"` + WorkspaceID string `db:"workspace_id"` + ForWorkspaceID sql.NullString `db:"for_workspace_id"` + Name sql.NullString `db:"name"` + OwnerID sql.NullString `db:"owner_id"` + IdentityID sql.NullString `db:"identity_id"` + Meta sql.NullString `db:"meta"` + Expires sql.NullTime `db:"expires"` + CreatedAtM int64 `db:"created_at_m"` + UpdatedAtM sql.NullInt64 `db:"updated_at_m"` + DeletedAtM sql.NullInt64 `db:"deleted_at_m"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + LastRefillAt sql.NullTime `db:"last_refill_at"` + Enabled bool `db:"enabled"` + RemainingRequests sql.NullInt32 `db:"remaining_requests"` + RatelimitAsync sql.NullBool `db:"ratelimit_async"` + RatelimitLimit sql.NullInt32 `db:"ratelimit_limit"` + RatelimitDuration sql.NullInt64 `db:"ratelimit_duration"` + Environment sql.NullString `db:"environment"` + Api Api `db:"api"` + KeyAuth KeyAuth `db:"key_auth"` + Workspace Workspace `db:"workspace"` + IdentityTableID sql.NullString `db:"identity_table_id"` + IdentityExternalID sql.NullString `db:"identity_external_id"` + IdentityMeta []byte `db:"identity_meta"` + EncryptedKey sql.NullString `db:"encrypted_key"` + EncryptionKeyID sql.NullString `db:"encryption_key_id"` + Roles interface{} `db:"roles"` + Permissions interface{} `db:"permissions"` + RolePermissions interface{} `db:"role_permissions"` + Ratelimits interface{} `db:"ratelimits"` + CreditID sql.NullString `db:"credit_id"` + CreditRemaining sql.NullInt32 `db:"credit_remaining"` + CreditRefillAmount sql.NullInt32 `db:"credit_refill_amount"` + CreditRefillDay sql.NullInt16 `db:"credit_refill_day"` + CreditRefilledAt sql.NullInt64 `db:"credit_refilled_at"` + IdentityCreditID sql.NullString `db:"identity_credit_id"` + IdentityCreditRemaining sql.NullInt32 `db:"identity_credit_remaining"` + IdentityCreditRefillAmount sql.NullInt32 `db:"identity_credit_refill_amount"` + IdentityCreditRefillDay sql.NullInt16 `db:"identity_credit_refill_day"` + IdentityCreditRefilledAt sql.NullInt64 `db:"identity_credit_refilled_at"` } // FindLiveKeyByHash @@ -216,7 +242,21 @@ type FindLiveKeyByHashRow struct { // FROM ratelimits rl // WHERE rl.key_id = k.id OR rl.identity_id = i.id), // JSON_ARRAY() -// ) as ratelimits +// ) as ratelimits, +// +// -- Key credits +// kc.id as credit_id, +// kc.remaining as credit_remaining, +// kc.refill_amount as credit_refill_amount, +// kc.refill_day as credit_refill_day, +// kc.refilled_at as credit_refilled_at, +// +// -- Identity credits +// ic.id as identity_credit_id, +// ic.remaining as identity_credit_remaining, +// ic.refill_amount as identity_credit_refill_amount, +// ic.refill_day as identity_credit_refill_day, +// ic.refilled_at as identity_credit_refilled_at // // FROM `keys` k // JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -224,6 +264,8 @@ type FindLiveKeyByHashRow struct { // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false // LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +// LEFT JOIN credits kc ON kc.key_id = k.id +// LEFT JOIN credits ic ON ic.identity_id = i.id // WHERE k.hash = ? // AND k.deleted_at_m IS NULL // AND a.deleted_at_m IS NULL @@ -302,6 +344,16 @@ func (q *Queries) FindLiveKeyByHash(ctx context.Context, db DBTX, hash string) ( &i.Permissions, &i.RolePermissions, &i.Ratelimits, + &i.CreditID, + &i.CreditRemaining, + &i.CreditRefillAmount, + &i.CreditRefillDay, + &i.CreditRefilledAt, + &i.IdentityCreditID, + &i.IdentityCreditRemaining, + &i.IdentityCreditRefillAmount, + &i.IdentityCreditRefillDay, + &i.IdentityCreditRefilledAt, ) return i, err } diff --git a/go/pkg/db/key_find_live_by_id.sql_generated.go b/go/pkg/db/key_find_live_by_id.sql_generated.go index db6a9d0b7e..f89eaa6e0c 100644 --- a/go/pkg/db/key_find_live_by_id.sql_generated.go +++ b/go/pkg/db/key_find_live_by_id.sql_generated.go @@ -87,7 +87,20 @@ SELECT WHERE rl.key_id = k.id OR rl.identity_id = i.id), JSON_ARRAY() - ) as ratelimits + ) as ratelimits, + + -- Credits Key/Identity based + kc.id as credit_id, + kc.remaining as credit_remaining, + kc.refill_amount as credit_refill_amount, + kc.refill_day as credit_refill_day, + kc.refilled_at as credit_refilled_at, + + ic.id as identity_credit_id, + ic.remaining as identity_credit_remaining, + ic.refill_amount as identity_credit_refill_amount, + ic.refill_day as identity_credit_refill_day, + ic.refilled_at as identity_credit_refilled_at FROM ` + "`" + `keys` + "`" + ` k JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -95,6 +108,8 @@ JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +LEFT JOIN credits kc ON kc.key_id = k.id +LEFT JOIN credits ic ON ic.identity_id = i.id WHERE k.id = ? AND k.deleted_at_m IS NULL AND a.deleted_at_m IS NULL @@ -103,41 +118,51 @@ WHERE k.id = ? ` type FindLiveKeyByIDRow struct { - ID string `db:"id"` - KeyAuthID string `db:"key_auth_id"` - Hash string `db:"hash"` - Start string `db:"start"` - WorkspaceID string `db:"workspace_id"` - ForWorkspaceID sql.NullString `db:"for_workspace_id"` - Name sql.NullString `db:"name"` - OwnerID sql.NullString `db:"owner_id"` - IdentityID sql.NullString `db:"identity_id"` - Meta sql.NullString `db:"meta"` - Expires sql.NullTime `db:"expires"` - CreatedAtM int64 `db:"created_at_m"` - UpdatedAtM sql.NullInt64 `db:"updated_at_m"` - DeletedAtM sql.NullInt64 `db:"deleted_at_m"` - RefillDay sql.NullInt16 `db:"refill_day"` - RefillAmount sql.NullInt32 `db:"refill_amount"` - LastRefillAt sql.NullTime `db:"last_refill_at"` - Enabled bool `db:"enabled"` - RemainingRequests sql.NullInt32 `db:"remaining_requests"` - RatelimitAsync sql.NullBool `db:"ratelimit_async"` - RatelimitLimit sql.NullInt32 `db:"ratelimit_limit"` - RatelimitDuration sql.NullInt64 `db:"ratelimit_duration"` - Environment sql.NullString `db:"environment"` - Api Api `db:"api"` - KeyAuth KeyAuth `db:"key_auth"` - Workspace Workspace `db:"workspace"` - IdentityTableID sql.NullString `db:"identity_table_id"` - IdentityExternalID sql.NullString `db:"identity_external_id"` - IdentityMeta []byte `db:"identity_meta"` - EncryptedKey sql.NullString `db:"encrypted_key"` - EncryptionKeyID sql.NullString `db:"encryption_key_id"` - Roles interface{} `db:"roles"` - Permissions interface{} `db:"permissions"` - RolePermissions interface{} `db:"role_permissions"` - Ratelimits interface{} `db:"ratelimits"` + ID string `db:"id"` + KeyAuthID string `db:"key_auth_id"` + Hash string `db:"hash"` + Start string `db:"start"` + WorkspaceID string `db:"workspace_id"` + ForWorkspaceID sql.NullString `db:"for_workspace_id"` + Name sql.NullString `db:"name"` + OwnerID sql.NullString `db:"owner_id"` + IdentityID sql.NullString `db:"identity_id"` + Meta sql.NullString `db:"meta"` + Expires sql.NullTime `db:"expires"` + CreatedAtM int64 `db:"created_at_m"` + UpdatedAtM sql.NullInt64 `db:"updated_at_m"` + DeletedAtM sql.NullInt64 `db:"deleted_at_m"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + LastRefillAt sql.NullTime `db:"last_refill_at"` + Enabled bool `db:"enabled"` + RemainingRequests sql.NullInt32 `db:"remaining_requests"` + RatelimitAsync sql.NullBool `db:"ratelimit_async"` + RatelimitLimit sql.NullInt32 `db:"ratelimit_limit"` + RatelimitDuration sql.NullInt64 `db:"ratelimit_duration"` + Environment sql.NullString `db:"environment"` + Api Api `db:"api"` + KeyAuth KeyAuth `db:"key_auth"` + Workspace Workspace `db:"workspace"` + IdentityTableID sql.NullString `db:"identity_table_id"` + IdentityExternalID sql.NullString `db:"identity_external_id"` + IdentityMeta []byte `db:"identity_meta"` + EncryptedKey sql.NullString `db:"encrypted_key"` + EncryptionKeyID sql.NullString `db:"encryption_key_id"` + Roles interface{} `db:"roles"` + Permissions interface{} `db:"permissions"` + RolePermissions interface{} `db:"role_permissions"` + Ratelimits interface{} `db:"ratelimits"` + CreditID sql.NullString `db:"credit_id"` + CreditRemaining sql.NullInt32 `db:"credit_remaining"` + CreditRefillAmount sql.NullInt32 `db:"credit_refill_amount"` + CreditRefillDay sql.NullInt16 `db:"credit_refill_day"` + CreditRefilledAt sql.NullInt64 `db:"credit_refilled_at"` + IdentityCreditID sql.NullString `db:"identity_credit_id"` + IdentityCreditRemaining sql.NullInt32 `db:"identity_credit_remaining"` + IdentityCreditRefillAmount sql.NullInt32 `db:"identity_credit_refill_amount"` + IdentityCreditRefillDay sql.NullInt16 `db:"identity_credit_refill_day"` + IdentityCreditRefilledAt sql.NullInt64 `db:"identity_credit_refilled_at"` } // FindLiveKeyByID @@ -218,7 +243,20 @@ type FindLiveKeyByIDRow struct { // WHERE rl.key_id = k.id // OR rl.identity_id = i.id), // JSON_ARRAY() -// ) as ratelimits +// ) as ratelimits, +// +// -- Credits Key/Identity based +// kc.id as credit_id, +// kc.remaining as credit_remaining, +// kc.refill_amount as credit_refill_amount, +// kc.refill_day as credit_refill_day, +// kc.refilled_at as credit_refilled_at, +// +// ic.id as identity_credit_id, +// ic.remaining as identity_credit_remaining, +// ic.refill_amount as identity_credit_refill_amount, +// ic.refill_day as identity_credit_refill_day, +// ic.refilled_at as identity_credit_refilled_at // // FROM `keys` k // JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -226,6 +264,8 @@ type FindLiveKeyByIDRow struct { // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false // LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +// LEFT JOIN credits kc ON kc.key_id = k.id +// LEFT JOIN credits ic ON ic.identity_id = i.id // WHERE k.id = ? // AND k.deleted_at_m IS NULL // AND a.deleted_at_m IS NULL @@ -304,6 +344,16 @@ func (q *Queries) FindLiveKeyByID(ctx context.Context, db DBTX, id string) (Find &i.Permissions, &i.RolePermissions, &i.Ratelimits, + &i.CreditID, + &i.CreditRemaining, + &i.CreditRefillAmount, + &i.CreditRefillDay, + &i.CreditRefilledAt, + &i.IdentityCreditID, + &i.IdentityCreditRemaining, + &i.IdentityCreditRefillAmount, + &i.IdentityCreditRefillDay, + &i.IdentityCreditRefilledAt, ) return i, err } diff --git a/go/pkg/db/key_find_remaining.sql_generated.go b/go/pkg/db/key_find_remaining.sql_generated.go new file mode 100644 index 0000000000..d31ed907d7 --- /dev/null +++ b/go/pkg/db/key_find_remaining.sql_generated.go @@ -0,0 +1,25 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: key_find_remaining.sql + +package db + +import ( + "context" + "database/sql" +) + +const findRemainingKey = `-- name: FindRemainingKey :one +SELECT remaining_requests FROM ` + "`" + `keys` + "`" + ` WHERE id = ? +` + +// FindRemainingKey +// +// SELECT remaining_requests FROM `keys` WHERE id = ? +func (q *Queries) FindRemainingKey(ctx context.Context, db DBTX, id string) (sql.NullInt32, error) { + row := db.QueryRowContext(ctx, findRemainingKey, id) + var remaining_requests sql.NullInt32 + err := row.Scan(&remaining_requests) + return remaining_requests, err +} diff --git a/go/pkg/db/keys_find_without_credits.sql_generated.go b/go/pkg/db/keys_find_without_credits.sql_generated.go new file mode 100644 index 0000000000..053ec83e3b --- /dev/null +++ b/go/pkg/db/keys_find_without_credits.sql_generated.go @@ -0,0 +1,104 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.29.0 +// source: keys_find_without_credits.sql + +package db + +import ( + "context" + "database/sql" +) + +const findKeysWithoutCredits = `-- name: FindKeysWithoutCredits :many +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 +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 + AND c.id IS NULL +ORDER BY k.created_at_m DESC +LIMIT ? +OFFSET ? +` + +type FindKeysWithoutCreditsParams struct { + Limit int32 `db:"limit"` + Offset int32 `db:"offset"` +} + +type FindKeysWithoutCreditsRow struct { + ID string `db:"id"` + WorkspaceID string `db:"workspace_id"` + 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"` + CreatedAtM int64 `db:"created_at_m"` + UpdatedAtM sql.NullInt64 `db:"updated_at_m"` +} + +// FindKeysWithoutCredits +// +// 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 +// 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 +// AND c.id IS NULL +// ORDER BY k.created_at_m DESC +// LIMIT ? +// OFFSET ? +func (q *Queries) FindKeysWithoutCredits(ctx context.Context, db DBTX, arg FindKeysWithoutCreditsParams) ([]FindKeysWithoutCreditsRow, error) { + rows, err := db.QueryContext(ctx, findKeysWithoutCredits, arg.Limit, arg.Offset) + if err != nil { + return nil, err + } + defer rows.Close() + var items []FindKeysWithoutCreditsRow + for rows.Next() { + var i FindKeysWithoutCreditsRow + if err := rows.Scan( + &i.ID, + &i.WorkspaceID, + &i.RemainingRequests, + &i.RefillDay, + &i.RefillAmount, + &i.LastRefillAtUnix, + &i.CreatedAtM, + &i.UpdatedAtM, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/go/pkg/db/models_generated.go b/go/pkg/db/models_generated.go index a37822ffdd..f220171468 100644 --- a/go/pkg/db/models_generated.go +++ b/go/pkg/db/models_generated.go @@ -580,6 +580,19 @@ type ClickhouseWorkspaceSetting struct { UpdatedAt sql.NullInt64 `db:"updated_at"` } +type Credit struct { + ID string `db:"id"` + WorkspaceID string `db:"workspace_id"` + KeyID sql.NullString `db:"key_id"` + IdentityID sql.NullString `db:"identity_id"` + Remaining int32 `db:"remaining"` + RefillDay sql.NullInt16 `db:"refill_day"` + RefillAmount sql.NullInt32 `db:"refill_amount"` + RefilledAt sql.NullInt64 `db:"refilled_at"` + CreatedAt int64 `db:"created_at"` + UpdatedAt sql.NullInt64 `db:"updated_at"` +} + type Deployment struct { ID string `db:"id"` WorkspaceID string `db:"workspace_id"` diff --git a/go/pkg/db/querier_bulk_generated.go b/go/pkg/db/querier_bulk_generated.go index 6e1be6a4f5..494ba9d874 100644 --- a/go/pkg/db/querier_bulk_generated.go +++ b/go/pkg/db/querier_bulk_generated.go @@ -12,6 +12,8 @@ type BulkQuerier interface { InsertAuditLogs(ctx context.Context, db DBTX, args []InsertAuditLogParams) error InsertAuditLogTargets(ctx context.Context, db DBTX, args []InsertAuditLogTargetParams) error InsertClickhouseWorkspaceSettingses(ctx context.Context, db DBTX, args []InsertClickhouseWorkspaceSettingsParams) error + InsertCredits(ctx context.Context, db DBTX, args []InsertCreditParams) error + UpsertCredit(ctx context.Context, db DBTX, args []UpsertCreditParams) error InsertDeployments(ctx context.Context, db DBTX, args []InsertDeploymentParams) error InsertDeploymentSteps(ctx context.Context, db DBTX, args []InsertDeploymentStepParams) error InsertDomains(ctx context.Context, db DBTX, args []InsertDomainParams) error diff --git a/go/pkg/db/querier_generated.go b/go/pkg/db/querier_generated.go index a19e3a22c8..ca1401c96f 100644 --- a/go/pkg/db/querier_generated.go +++ b/go/pkg/db/querier_generated.go @@ -26,6 +26,10 @@ type Querier interface { // DELETE FROM keys_roles // WHERE key_id = ? DeleteAllKeyRolesByKeyID(ctx context.Context, db DBTX, keyID string) error + //DeleteCredit + // + // DELETE FROM credits WHERE id = ? + DeleteCredit(ctx context.Context, db DBTX, id string) error //DeleteIdentity // // DELETE FROM identities @@ -92,9 +96,10 @@ type Querier interface { DeleteManyRolePermissionsByRoleID(ctx context.Context, db DBTX, roleID string) error //DeleteOldIdentityByExternalID // - // DELETE i, rl + // DELETE i, rl, c // FROM identities i // LEFT JOIN ratelimits rl ON rl.identity_id = i.id + // LEFT JOIN credits c ON c.identity_id = i.id // WHERE i.workspace_id = ? // AND i.external_id = ? // AND i.id != ? @@ -102,9 +107,10 @@ type Querier interface { DeleteOldIdentityByExternalID(ctx context.Context, db DBTX, arg DeleteOldIdentityByExternalIDParams) error //DeleteOldIdentityWithRatelimits // - // DELETE i, rl + // DELETE i, rl, c // FROM identities i // LEFT JOIN ratelimits rl ON rl.identity_id = i.id + // LEFT JOIN credits c ON c.identity_id = i.id // WHERE i.workspace_id = ? // AND (i.id = ? OR i.external_id = ?) // AND i.deleted = true @@ -153,6 +159,14 @@ type Querier interface { // SELECT workspace_id, username, password_encrypted, quota_duration_seconds, max_queries_per_window, max_execution_time_per_window, max_query_execution_time, max_query_memory_bytes, max_query_result_rows, created_at, updated_at FROM `clickhouse_workspace_settings` // WHERE workspace_id = ? FindClickhouseWorkspaceSettingsByWorkspaceID(ctx context.Context, db DBTX, workspaceID string) (ClickhouseWorkspaceSetting, error) + //FindCreditsByIdentityID + // + // SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM `credits` WHERE identity_id = ? + FindCreditsByIdentityID(ctx context.Context, db DBTX, identityID sql.NullString) (Credit, error) + //FindCreditsByKeyID + // + // SELECT id, workspace_id, key_id, identity_id, remaining, refill_day, refill_amount, refilled_at, created_at, updated_at FROM `credits` WHERE key_id = ? + FindCreditsByKeyID(ctx context.Context, db DBTX, keyID sql.NullString) (Credit, error) //FindDeploymentById // // SELECT @@ -278,26 +292,15 @@ type Querier interface { // AND deleted = ? // AND (external_id IN(/*SLICE:identities*/?) OR id IN (/*SLICE:identities*/?)) FindIdentities(ctx context.Context, db DBTX, arg FindIdentitiesParams) ([]Identity, error) - //FindIdentityByExternalID - // - // SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at - // FROM identities - // WHERE workspace_id = ? - // AND external_id = ? - // AND deleted = ? - FindIdentityByExternalID(ctx context.Context, db DBTX, arg FindIdentityByExternalIDParams) (Identity, error) - //FindIdentityByID - // - // SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at - // FROM identities - // WHERE workspace_id = ? - // AND id = ? - // AND deleted = ? - FindIdentityByID(ctx context.Context, db DBTX, arg FindIdentityByIDParams) (Identity, error) - //FindIdentityWithRatelimits + //FindIdentity // // SELECT // i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, + // c.id AS credit_id, + // c.remaining AS credit_remaining, + // c.refill_amount AS credit_refill_amount, + // c.refill_day AS credit_refill_day, + // c.refilled_at AS credit_refilled_at, // COALESCE( // (SELECT JSON_ARRAYAGG( // JSON_OBJECT( @@ -314,33 +317,36 @@ type Querier interface { // JSON_ARRAY() // ) as ratelimits // FROM identities i - // WHERE i.workspace_id = ? - // AND i.id = ? - // AND i.deleted = ? - // UNION ALL - // SELECT - // i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, - // COALESCE( - // (SELECT JSON_ARRAYAGG( - // JSON_OBJECT( - // 'id', rl.id, - // 'name', rl.name, - // 'key_id', rl.key_id, - // 'identity_id', rl.identity_id, - // 'limit', rl.`limit`, - // 'duration', rl.duration, - // 'auto_apply', rl.auto_apply = 1 - // ) - // ) - // FROM ratelimits rl WHERE rl.identity_id = i.id), - // JSON_ARRAY() - // ) as ratelimits - // FROM identities i - // WHERE i.workspace_id = ? - // AND i.external_id = ? - // AND i.deleted = ? - // LIMIT 1 - FindIdentityWithRatelimits(ctx context.Context, db DBTX, arg FindIdentityWithRatelimitsParams) ([]FindIdentityWithRatelimitsRow, error) + // LEFT JOIN credits c ON c.identity_id = i.id + // WHERE i.id = ( + // SELECT id FROM identities sub1 + // WHERE sub1.workspace_id = ? + // AND sub1.id = ? + // AND sub1.deleted = ? + // UNION ALL + // SELECT id FROM identities sub2 + // WHERE sub2.workspace_id = ? + // AND sub2.external_id = ? + // AND sub2.deleted = ? + // LIMIT 1 + // ) + FindIdentity(ctx context.Context, db DBTX, arg FindIdentityParams) (FindIdentityRow, error) + //FindIdentityByExternalID + // + // SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at + // FROM identities + // WHERE workspace_id = ? + // AND external_id = ? + // AND deleted = ? + FindIdentityByExternalID(ctx context.Context, db DBTX, arg FindIdentityByExternalIDParams) (Identity, error) + //FindIdentityByID + // + // SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at + // FROM identities + // WHERE workspace_id = ? + // AND id = ? + // AND deleted = ? + FindIdentityByID(ctx context.Context, db DBTX, arg FindIdentityByIDParams) (Identity, error) //FindKeyAuthsByIds // // SELECT ka.id as key_auth_id, a.id as api_id @@ -443,13 +449,30 @@ type Querier interface { // i.meta as identity_meta, // ka.deleted_at_m as key_auth_deleted_at_m, // ws.enabled as workspace_enabled, - // fws.enabled as for_workspace_enabled + // fws.enabled as for_workspace_enabled, + // + // -- Key-level credits + // key_credits.id as key_credit_id, + // key_credits.remaining as key_credit_remaining, + // key_credits.refill_amount as key_credit_refill_amount, + // key_credits.refill_day as key_credit_refill_day, + // key_credits.refilled_at as key_credit_refilled_at, + // + // -- Identity-level credits + // identity_credits.id as identity_credit_id, + // identity_credits.remaining as identity_credit_remaining, + // identity_credits.refill_amount as identity_credit_refill_amount, + // identity_credits.refill_day as identity_credit_refill_day, + // identity_credits.refilled_at as identity_credit_refilled_at + // // from `keys` k // JOIN apis a USING (key_auth_id) // JOIN key_auth ka ON ka.id = k.key_auth_id // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN workspaces fws ON fws.id = k.for_workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = 0 + // LEFT JOIN credits key_credits ON key_credits.key_id = k.id + // LEFT JOIN credits identity_credits ON identity_credits.identity_id = i.id // where k.hash = ? // and k.deleted_at_m is null FindKeyForVerification(ctx context.Context, db DBTX, hash string) (FindKeyForVerificationRow, error) @@ -464,6 +487,29 @@ type Querier interface { // // SELECT id, workspace_id, created_at_m, updated_at_m, deleted_at_m, store_encrypted_keys, default_prefix, default_bytes, size_approx, size_last_updated_at FROM `key_auth` WHERE id = ? FindKeySpaceByID(ctx context.Context, db DBTX, id string) (KeyAuth, error) + //FindKeysWithoutCredits + // + // 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 + // 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 + // AND c.id IS NULL + // ORDER BY k.created_at_m DESC + // LIMIT ? + // OFFSET ? + FindKeysWithoutCredits(ctx context.Context, db DBTX, arg FindKeysWithoutCreditsParams) ([]FindKeysWithoutCreditsRow, error) //FindLiveApiByID // // SELECT apis.id, apis.name, apis.workspace_id, apis.ip_whitelist, apis.auth_type, apis.key_auth_id, apis.created_at_m, apis.updated_at_m, apis.deleted_at_m, apis.delete_protection, ka.id, ka.workspace_id, ka.created_at_m, ka.updated_at_m, ka.deleted_at_m, ka.store_encrypted_keys, ka.default_prefix, ka.default_bytes, ka.size_approx, ka.size_last_updated_at @@ -551,7 +597,21 @@ type Querier interface { // FROM ratelimits rl // WHERE rl.key_id = k.id OR rl.identity_id = i.id), // JSON_ARRAY() - // ) as ratelimits + // ) as ratelimits, + // + // -- Key credits + // kc.id as credit_id, + // kc.remaining as credit_remaining, + // kc.refill_amount as credit_refill_amount, + // kc.refill_day as credit_refill_day, + // kc.refilled_at as credit_refilled_at, + // + // -- Identity credits + // ic.id as identity_credit_id, + // ic.remaining as identity_credit_remaining, + // ic.refill_amount as identity_credit_refill_amount, + // ic.refill_day as identity_credit_refill_day, + // ic.refilled_at as identity_credit_refilled_at // // FROM `keys` k // JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -559,6 +619,8 @@ type Querier interface { // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false // LEFT JOIN encrypted_keys ek ON ek.key_id = k.id + // LEFT JOIN credits kc ON kc.key_id = k.id + // LEFT JOIN credits ic ON ic.identity_id = i.id // WHERE k.hash = ? // AND k.deleted_at_m IS NULL // AND a.deleted_at_m IS NULL @@ -643,7 +705,20 @@ type Querier interface { // WHERE rl.key_id = k.id // OR rl.identity_id = i.id), // JSON_ARRAY() - // ) as ratelimits + // ) as ratelimits, + // + // -- Credits Key/Identity based + // kc.id as credit_id, + // kc.remaining as credit_remaining, + // kc.refill_amount as credit_refill_amount, + // kc.refill_day as credit_refill_day, + // kc.refilled_at as credit_refilled_at, + // + // ic.id as identity_credit_id, + // ic.remaining as identity_credit_remaining, + // ic.refill_amount as identity_credit_refill_amount, + // ic.refill_day as identity_credit_refill_day, + // ic.refilled_at as identity_credit_refilled_at // // FROM `keys` k // JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -651,6 +726,8 @@ type Querier interface { // JOIN workspaces ws ON ws.id = k.workspace_id // LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false // LEFT JOIN encrypted_keys ek ON ek.key_id = k.id + // LEFT JOIN credits kc ON kc.key_id = k.id + // LEFT JOIN credits ic ON ic.identity_id = i.id // WHERE k.id = ? // AND k.deleted_at_m IS NULL // AND a.deleted_at_m IS NULL @@ -813,6 +890,14 @@ type Querier interface { // AND namespace_id = ? // AND identifier = ? FindRatelimitOverrideByIdentifier(ctx context.Context, db DBTX, arg FindRatelimitOverrideByIdentifierParams) (RatelimitOverride, error) + //FindRemainingCredits + // + // SELECT remaining FROM `credits` WHERE id = ? + FindRemainingCredits(ctx context.Context, db DBTX, id string) (int32, error) + //FindRemainingKey + // + // SELECT remaining_requests FROM `keys` WHERE id = ? + FindRemainingKey(ctx context.Context, db DBTX, id string) (sql.NullInt32, error) // Finds a role record by its ID // Returns: The role record if found // @@ -1029,6 +1114,21 @@ type Querier interface { // ? // ) InsertClickhouseWorkspaceSettings(ctx context.Context, db DBTX, arg InsertClickhouseWorkspaceSettingsParams) error + //InsertCredit + // + // INSERT INTO `credits` ( + // id, + // workspace_id, + // key_id, + // identity_id, + // remaining, + // refill_day, + // refill_amount, + // created_at, + // updated_at, + // refilled_at + // ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + InsertCredit(ctx context.Context, db DBTX, arg InsertCreditParams) error //InsertDeployment // // INSERT INTO `deployments` ( @@ -1458,14 +1558,20 @@ type Querier interface { ListExecutableChallenges(ctx context.Context, db DBTX, verificationTypes []AcmeChallengesType) ([]ListExecutableChallengesRow, error) //ListIdentities // - // SELECT id, external_id, workspace_id, environment, meta, deleted, created_at, updated_at - // FROM identities - // WHERE workspace_id = ? - // AND deleted = ? - // AND id >= ? - // ORDER BY id ASC + // SELECT + // i.id, i.external_id, i.workspace_id, i.environment, i.meta, i.deleted, i.created_at, i.updated_at, + // c.id as credit_id, + // c.remaining as credit_remaining, + // c.refill_amount as credit_refill_amount, + // c.refill_day as credit_refill_day + // FROM identities i + // LEFT JOIN credits c ON c.identity_id = i.id + // WHERE i.workspace_id = ? + // AND i.deleted = ? + // AND i.id >= ? + // ORDER BY i.id ASC // LIMIT ? - ListIdentities(ctx context.Context, db DBTX, arg ListIdentitiesParams) ([]Identity, error) + ListIdentities(ctx context.Context, db DBTX, arg ListIdentitiesParams) ([]ListIdentitiesRow, error) //ListIdentityRatelimits // // SELECT id, name, workspace_id, created_at, updated_at, key_id, identity_id, `limit`, duration, auto_apply @@ -1828,6 +1934,30 @@ type Querier interface { // updated_at = ? // WHERE workspace_id = ? UpdateClickhouseWorkspaceSettingsLimits(ctx context.Context, db DBTX, arg UpdateClickhouseWorkspaceSettingsLimitsParams) error + //UpdateCreditDecrement + // + // UPDATE `credits` + // SET remaining = CASE + // WHEN remaining >= ? THEN remaining - ? + // ELSE 0 + // END, + // updated_at = ? + // WHERE id = ? + UpdateCreditDecrement(ctx context.Context, db DBTX, arg UpdateCreditDecrementParams) error + //UpdateCreditIncrement + // + // UPDATE `credits` + // SET remaining = remaining + ?, + // updated_at = ? + // WHERE id = ? + UpdateCreditIncrement(ctx context.Context, db DBTX, arg UpdateCreditIncrementParams) error + //UpdateCreditSet + // + // UPDATE `credits` + // SET remaining = ?, + // updated_at = ? + // WHERE id = ? + UpdateCreditSet(ctx context.Context, db DBTX, arg UpdateCreditSetParams) error //UpdateDeploymentOpenapiSpec // // UPDATE deployments @@ -1967,6 +2097,39 @@ type Querier interface { // SET plan = ? // WHERE id = ? UpdateWorkspacePlan(ctx context.Context, db DBTX, arg UpdateWorkspacePlanParams) (sql.Result, error) + //UpsertCredit + // + // INSERT INTO `credits` ( + // id, + // workspace_id, + // key_id, + // identity_id, + // remaining, + // refill_day, + // refill_amount, + // created_at, + // updated_at, + // refilled_at + // ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + // ON DUPLICATE KEY UPDATE + // remaining = CASE + // WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(remaining) + // ELSE remaining + // END, + // refill_day = CASE + // WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_day) + // ELSE refill_day + // END, + // refill_amount = CASE + // WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refill_amount) + // ELSE refill_amount + // END, + // refilled_at = CASE + // WHEN CAST(? AS UNSIGNED) = 1 THEN VALUES(refilled_at) + // ELSE refilled_at + // END, + // updated_at = VALUES(updated_at) + UpsertCredit(ctx context.Context, db DBTX, arg UpsertCreditParams) error } var _ Querier = (*Queries)(nil) diff --git a/go/pkg/db/queries/credits_delete.sql b/go/pkg/db/queries/credits_delete.sql new file mode 100644 index 0000000000..3516d5eb44 --- /dev/null +++ b/go/pkg/db/queries/credits_delete.sql @@ -0,0 +1,2 @@ +-- name: DeleteCredit :exec +DELETE FROM credits WHERE id = ?; diff --git a/go/pkg/db/queries/credits_find_by_identity_id.sql b/go/pkg/db/queries/credits_find_by_identity_id.sql new file mode 100644 index 0000000000..46fbb27031 --- /dev/null +++ b/go/pkg/db/queries/credits_find_by_identity_id.sql @@ -0,0 +1,2 @@ +-- name: FindCreditsByIdentityID :one +SELECT * FROM `credits` WHERE identity_id = ?; diff --git a/go/pkg/db/queries/credits_find_by_key_id.sql b/go/pkg/db/queries/credits_find_by_key_id.sql new file mode 100644 index 0000000000..4a28bb88d4 --- /dev/null +++ b/go/pkg/db/queries/credits_find_by_key_id.sql @@ -0,0 +1,2 @@ +-- name: FindCreditsByKeyID :one +SELECT * FROM `credits` WHERE key_id = ?; diff --git a/go/pkg/db/queries/credits_find_remaining.sql b/go/pkg/db/queries/credits_find_remaining.sql new file mode 100644 index 0000000000..0248a41412 --- /dev/null +++ b/go/pkg/db/queries/credits_find_remaining.sql @@ -0,0 +1,2 @@ +-- name: FindRemainingCredits :one +SELECT remaining FROM `credits` WHERE id = ?; diff --git a/go/pkg/db/queries/credits_insert.sql b/go/pkg/db/queries/credits_insert.sql new file mode 100644 index 0000000000..cfa929fd8c --- /dev/null +++ b/go/pkg/db/queries/credits_insert.sql @@ -0,0 +1,13 @@ +-- name: InsertCredit :exec +INSERT INTO `credits` ( + id, + workspace_id, + key_id, + identity_id, + remaining, + refill_day, + refill_amount, + created_at, + updated_at, + refilled_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?); diff --git a/go/pkg/db/queries/credits_update_decrement.sql b/go/pkg/db/queries/credits_update_decrement.sql new file mode 100644 index 0000000000..2b0a547630 --- /dev/null +++ b/go/pkg/db/queries/credits_update_decrement.sql @@ -0,0 +1,8 @@ +-- name: UpdateCreditDecrement :exec +UPDATE `credits` +SET remaining = CASE + WHEN remaining >= sqlc.arg('credits') THEN remaining - sqlc.arg('credits') + ELSE 0 +END, + updated_at = ? +WHERE id = sqlc.arg('id'); diff --git a/go/pkg/db/queries/credits_update_increment.sql b/go/pkg/db/queries/credits_update_increment.sql new file mode 100644 index 0000000000..4cf7e2d98b --- /dev/null +++ b/go/pkg/db/queries/credits_update_increment.sql @@ -0,0 +1,5 @@ +-- name: UpdateCreditIncrement :exec +UPDATE `credits` +SET remaining = remaining + sqlc.arg('credits'), + updated_at = ? +WHERE id = ?; diff --git a/go/pkg/db/queries/credits_update_set.sql b/go/pkg/db/queries/credits_update_set.sql new file mode 100644 index 0000000000..9ba0595b68 --- /dev/null +++ b/go/pkg/db/queries/credits_update_set.sql @@ -0,0 +1,5 @@ +-- name: UpdateCreditSet :exec +UPDATE `credits` +SET remaining = sqlc.arg('remaining'), + updated_at = ? +WHERE id = ?; diff --git a/go/pkg/db/queries/credits_upsert.sql b/go/pkg/db/queries/credits_upsert.sql new file mode 100644 index 0000000000..6ce7ef6baa --- /dev/null +++ b/go/pkg/db/queries/credits_upsert.sql @@ -0,0 +1,31 @@ +-- name: UpsertCredit :exec +INSERT INTO `credits` ( + id, + workspace_id, + key_id, + identity_id, + remaining, + refill_day, + refill_amount, + created_at, + updated_at, + refilled_at +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +ON DUPLICATE KEY UPDATE + remaining = CASE + WHEN CAST(sqlc.arg('remaining_specified') AS UNSIGNED) = 1 THEN VALUES(remaining) + ELSE remaining + END, + refill_day = CASE + WHEN CAST(sqlc.arg('refill_day_specified') AS UNSIGNED) = 1 THEN VALUES(refill_day) + ELSE refill_day + END, + refill_amount = CASE + WHEN CAST(sqlc.arg('refill_amount_specified') AS UNSIGNED) = 1 THEN VALUES(refill_amount) + ELSE refill_amount + END, + refilled_at = CASE + WHEN CAST(sqlc.arg('refilled_at_specified') AS UNSIGNED) = 1 THEN VALUES(refilled_at) + ELSE refilled_at + END, + updated_at = VALUES(updated_at); diff --git a/go/pkg/db/queries/identity_delete_old_by_external_id.sql b/go/pkg/db/queries/identity_delete_old_by_external_id.sql index 070ffd2bc1..1cd0d2deee 100644 --- a/go/pkg/db/queries/identity_delete_old_by_external_id.sql +++ b/go/pkg/db/queries/identity_delete_old_by_external_id.sql @@ -1,7 +1,8 @@ -- name: DeleteOldIdentityByExternalID :exec -DELETE i, rl +DELETE i, rl, c FROM identities i LEFT JOIN ratelimits rl ON rl.identity_id = i.id +LEFT JOIN credits c ON c.identity_id = i.id WHERE i.workspace_id = sqlc.arg(workspace_id) AND i.external_id = sqlc.arg(external_id) AND i.id != sqlc.arg(current_identity_id) diff --git a/go/pkg/db/queries/identity_delete_old_with_ratelimits.sql b/go/pkg/db/queries/identity_delete_old_with_ratelimits.sql index 0aafb69366..20cbb73509 100644 --- a/go/pkg/db/queries/identity_delete_old_with_ratelimits.sql +++ b/go/pkg/db/queries/identity_delete_old_with_ratelimits.sql @@ -1,7 +1,8 @@ -- name: DeleteOldIdentityWithRatelimits :exec -DELETE i, rl +DELETE i, rl, c FROM identities i LEFT JOIN ratelimits rl ON rl.identity_id = i.id +LEFT JOIN credits c ON c.identity_id = i.id WHERE i.workspace_id = sqlc.arg(workspace_id) AND (i.id = sqlc.arg(identity) OR i.external_id = sqlc.arg(identity)) AND i.deleted = true; diff --git a/go/pkg/db/queries/identity_find.sql b/go/pkg/db/queries/identity_find.sql new file mode 100644 index 0000000000..9fad98a58c --- /dev/null +++ b/go/pkg/db/queries/identity_find.sql @@ -0,0 +1,37 @@ +-- name: FindIdentity :one +SELECT + i.*, + c.id AS credit_id, + c.remaining AS credit_remaining, + c.refill_amount AS credit_refill_amount, + c.refill_day AS credit_refill_day, + c.refilled_at AS credit_refilled_at, + COALESCE( + (SELECT JSON_ARRAYAGG( + JSON_OBJECT( + 'id', rl.id, + 'name', rl.name, + 'key_id', rl.key_id, + 'identity_id', rl.identity_id, + 'limit', rl.`limit`, + 'duration', rl.duration, + 'auto_apply', rl.auto_apply = 1 + ) + ) + FROM ratelimits rl WHERE rl.identity_id = i.id), + JSON_ARRAY() + ) as ratelimits +FROM identities i +LEFT JOIN credits c ON c.identity_id = i.id +WHERE i.id = ( + SELECT id FROM identities sub1 + WHERE sub1.workspace_id = sqlc.arg(workspace_id) + AND sub1.id = sqlc.arg(identity) + AND sub1.deleted = sqlc.arg(deleted) + UNION ALL + SELECT id FROM identities sub2 + WHERE sub2.workspace_id = sqlc.arg(workspace_id) + AND sub2.external_id = sqlc.arg(identity) + AND sub2.deleted = sqlc.arg(deleted) + LIMIT 1 +); diff --git a/go/pkg/db/queries/identity_find_with_ratelimits.sql b/go/pkg/db/queries/identity_find_with_ratelimits.sql deleted file mode 100644 index 9f7f24d457..0000000000 --- a/go/pkg/db/queries/identity_find_with_ratelimits.sql +++ /dev/null @@ -1,45 +0,0 @@ --- name: FindIdentityWithRatelimits :many -SELECT - i.*, - COALESCE( - (SELECT JSON_ARRAYAGG( - JSON_OBJECT( - 'id', rl.id, - 'name', rl.name, - 'key_id', rl.key_id, - 'identity_id', rl.identity_id, - 'limit', rl.`limit`, - 'duration', rl.duration, - 'auto_apply', rl.auto_apply = 1 - ) - ) - FROM ratelimits rl WHERE rl.identity_id = i.id), - JSON_ARRAY() - ) as ratelimits -FROM identities i -WHERE i.workspace_id = sqlc.arg(workspace_id) - AND i.id = sqlc.arg(identity) - AND i.deleted = sqlc.arg(deleted) -UNION ALL -SELECT - i.*, - COALESCE( - (SELECT JSON_ARRAYAGG( - JSON_OBJECT( - 'id', rl.id, - 'name', rl.name, - 'key_id', rl.key_id, - 'identity_id', rl.identity_id, - 'limit', rl.`limit`, - 'duration', rl.duration, - 'auto_apply', rl.auto_apply = 1 - ) - ) - FROM ratelimits rl WHERE rl.identity_id = i.id), - JSON_ARRAY() - ) as ratelimits -FROM identities i -WHERE i.workspace_id = sqlc.arg(workspace_id) - AND i.external_id = sqlc.arg(identity) - AND i.deleted = sqlc.arg(deleted) -LIMIT 1; diff --git a/go/pkg/db/queries/identity_list.sql b/go/pkg/db/queries/identity_list.sql index 0302e082de..6ae5941e55 100644 --- a/go/pkg/db/queries/identity_list.sql +++ b/go/pkg/db/queries/identity_list.sql @@ -1,8 +1,14 @@ -- name: ListIdentities :many -SELECT * -FROM identities -WHERE workspace_id = sqlc.arg(workspace_id) -AND deleted = sqlc.arg(deleted) -AND id >= sqlc.arg(id_cursor) -ORDER BY id ASC +SELECT + i.*, + c.id as credit_id, + c.remaining as credit_remaining, + c.refill_amount as credit_refill_amount, + c.refill_day as credit_refill_day +FROM identities i +LEFT JOIN credits c ON c.identity_id = i.id +WHERE i.workspace_id = sqlc.arg(workspace_id) +AND i.deleted = sqlc.arg(deleted) +AND i.id >= sqlc.arg(id_cursor) +ORDER BY i.id ASC LIMIT ? diff --git a/go/pkg/db/queries/key_find_for_verification.sql b/go/pkg/db/queries/key_find_for_verification.sql index d649de3567..8accee46c7 100644 --- a/go/pkg/db/queries/key_find_for_verification.sql +++ b/go/pkg/db/queries/key_find_for_verification.sql @@ -66,12 +66,29 @@ select k.id, i.meta as identity_meta, ka.deleted_at_m as key_auth_deleted_at_m, ws.enabled as workspace_enabled, - fws.enabled as for_workspace_enabled + fws.enabled as for_workspace_enabled, + + -- Key-level credits + key_credits.id as key_credit_id, + key_credits.remaining as key_credit_remaining, + key_credits.refill_amount as key_credit_refill_amount, + key_credits.refill_day as key_credit_refill_day, + key_credits.refilled_at as key_credit_refilled_at, + + -- Identity-level credits + identity_credits.id as identity_credit_id, + identity_credits.remaining as identity_credit_remaining, + identity_credits.refill_amount as identity_credit_refill_amount, + identity_credits.refill_day as identity_credit_refill_day, + identity_credits.refilled_at as identity_credit_refilled_at + from `keys` k JOIN apis a USING (key_auth_id) JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN workspaces fws ON fws.id = k.for_workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = 0 + LEFT JOIN credits key_credits ON key_credits.key_id = k.id + LEFT JOIN credits identity_credits ON identity_credits.identity_id = i.id where k.hash = ? and k.deleted_at_m is null; diff --git a/go/pkg/db/queries/key_find_live_by_hash.sql b/go/pkg/db/queries/key_find_live_by_hash.sql index 71dbf64bbe..0833fe558f 100644 --- a/go/pkg/db/queries/key_find_live_by_hash.sql +++ b/go/pkg/db/queries/key_find_live_by_hash.sql @@ -74,7 +74,21 @@ SELECT FROM ratelimits rl WHERE rl.key_id = k.id OR rl.identity_id = i.id), JSON_ARRAY() - ) as ratelimits + ) as ratelimits, + + -- Key credits + kc.id as credit_id, + kc.remaining as credit_remaining, + kc.refill_amount as credit_refill_amount, + kc.refill_day as credit_refill_day, + kc.refilled_at as credit_refilled_at, + + -- Identity credits + ic.id as identity_credit_id, + ic.remaining as identity_credit_remaining, + ic.refill_amount as identity_credit_refill_amount, + ic.refill_day as identity_credit_refill_day, + ic.refilled_at as identity_credit_refilled_at FROM `keys` k JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -82,6 +96,8 @@ JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +LEFT JOIN credits kc ON kc.key_id = k.id +LEFT JOIN credits ic ON ic.identity_id = i.id WHERE k.hash = sqlc.arg(hash) AND k.deleted_at_m IS NULL AND a.deleted_at_m IS NULL diff --git a/go/pkg/db/queries/key_find_live_by_id.sql b/go/pkg/db/queries/key_find_live_by_id.sql index e76e900788..14a42bbd56 100644 --- a/go/pkg/db/queries/key_find_live_by_id.sql +++ b/go/pkg/db/queries/key_find_live_by_id.sql @@ -75,7 +75,20 @@ SELECT WHERE rl.key_id = k.id OR rl.identity_id = i.id), JSON_ARRAY() - ) as ratelimits + ) as ratelimits, + + -- Credits Key/Identity based + kc.id as credit_id, + kc.remaining as credit_remaining, + kc.refill_amount as credit_refill_amount, + kc.refill_day as credit_refill_day, + kc.refilled_at as credit_refilled_at, + + ic.id as identity_credit_id, + ic.remaining as identity_credit_remaining, + ic.refill_amount as identity_credit_refill_amount, + ic.refill_day as identity_credit_refill_day, + ic.refilled_at as identity_credit_refilled_at FROM `keys` k JOIN apis a ON a.key_auth_id = k.key_auth_id @@ -83,6 +96,8 @@ JOIN key_auth ka ON ka.id = k.key_auth_id JOIN workspaces ws ON ws.id = k.workspace_id LEFT JOIN identities i ON k.identity_id = i.id AND i.deleted = false LEFT JOIN encrypted_keys ek ON ek.key_id = k.id +LEFT JOIN credits kc ON kc.key_id = k.id +LEFT JOIN credits ic ON ic.identity_id = i.id WHERE k.id = sqlc.arg(id) AND k.deleted_at_m IS NULL AND a.deleted_at_m IS NULL diff --git a/go/pkg/db/queries/key_find_remaining.sql b/go/pkg/db/queries/key_find_remaining.sql new file mode 100644 index 0000000000..08f1c3d79e --- /dev/null +++ b/go/pkg/db/queries/key_find_remaining.sql @@ -0,0 +1,2 @@ +-- name: FindRemainingKey :one +SELECT remaining_requests FROM `keys` WHERE id = ?; diff --git a/go/pkg/db/queries/keys_find_without_credits.sql b/go/pkg/db/queries/keys_find_without_credits.sql new file mode 100644 index 0000000000..532916c41d --- /dev/null +++ b/go/pkg/db/queries/keys_find_without_credits.sql @@ -0,0 +1,21 @@ +-- name: FindKeysWithoutCredits :many +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 +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 + AND c.id IS NULL +ORDER BY k.created_at_m DESC +LIMIT ? +OFFSET ?; \ No newline at end of file diff --git a/go/pkg/db/schema.sql b/go/pkg/db/schema.sql index 0b6162b57c..4bf7c34a53 100644 --- a/go/pkg/db/schema.sql +++ b/go/pkg/db/schema.sql @@ -123,6 +123,22 @@ CREATE TABLE `keys` ( CONSTRAINT `hash_idx` UNIQUE(`hash`) ); +CREATE TABLE `credits` ( + `id` varchar(256) NOT NULL, + `workspace_id` varchar(256) NOT NULL, + `key_id` varchar(256), + `identity_id` varchar(256), + `remaining` int NOT NULL, + `refill_day` tinyint, + `refill_amount` int, + `refilled_at` bigint unsigned, + `created_at` bigint NOT NULL DEFAULT 0, + `updated_at` bigint, + CONSTRAINT `credits_id` PRIMARY KEY(`id`), + CONSTRAINT `unique_per_key_idx` UNIQUE(`key_id`), + CONSTRAINT `unique_per_identity_idx` UNIQUE(`identity_id`) +); + CREATE TABLE `vercel_bindings` ( `id` varchar(256) NOT NULL, `integration_id` varchar(256) NOT NULL, @@ -415,6 +431,8 @@ CREATE INDEX `idx_keys_on_workspace_id` ON `keys` (`workspace_id`); CREATE INDEX `owner_id_idx` ON `keys` (`owner_id`); CREATE INDEX `identity_id_idx` ON `keys` (`identity_id`); CREATE INDEX `deleted_at_idx` ON `keys` (`deleted_at_m`); +CREATE INDEX `workspace_id_idx` ON `credits` (`workspace_id`); +CREATE INDEX `refill_lookup_idx` ON `credits` (`refill_day`,`refill_amount`,`refilled_at`); CREATE INDEX `name_idx` ON `ratelimits` (`name`); CREATE INDEX `workspace_id_idx` ON `audit_log` (`workspace_id`); CREATE INDEX `bucket_id_idx` ON `audit_log` (`bucket_id`); diff --git a/go/pkg/uid/uid.go b/go/pkg/uid/uid.go index e3d0d14722..aae749ce29 100644 --- a/go/pkg/uid/uid.go +++ b/go/pkg/uid/uid.go @@ -27,6 +27,7 @@ const ( PermissionPrefix Prefix = "perm" IdentityPrefix Prefix = "id" RatelimitPrefix Prefix = "rl" + CreditPrefix Prefix = "cr" AuditLogBucketPrefix Prefix = "buk" AuditLogPrefix Prefix = "log" InstancePrefix Prefix = "ins" diff --git a/internal/db/src/schema/credits.ts b/internal/db/src/schema/credits.ts new file mode 100644 index 0000000000..d05427cf73 --- /dev/null +++ b/internal/db/src/schema/credits.ts @@ -0,0 +1,79 @@ +import { relations } from "drizzle-orm"; +import { + bigint, + index, + int, + mysqlTable, + tinyint, + uniqueIndex, + varchar, +} from "drizzle-orm/mysql-core"; +import { identities } from "./identity"; +import { keys } from "./keys"; +import { lifecycleDatesV2 } from "./util/lifecycle_dates"; +import { workspaces } from "./workspaces"; + +export const credits = mysqlTable( + "credits", + { + id: varchar("id", { length: 256 }).primaryKey(), + workspaceId: varchar("workspace_id", { length: 256 }).notNull(), + + /** + * Either keyId or identityId must be defined, not both + */ + keyId: varchar("key_id", { length: 256 }), + /** + * Either keyId or identityId must be defined, not both + */ + identityId: varchar("identity_id", { length: 256 }), + + /** + * The number of credits remaining + */ + remaining: int("remaining").notNull(), + + /** + * You can refill credits at a desired interval + * + * Specify the day on which we should refill. + * - 1 = we refill on the first of the month + * - 2 = we refill on the 2nd of the month + * - 31 = we refill on the 31st or last available day + * - null = we refill every day + */ + refillDay: tinyint("refill_day"), + refillAmount: int("refill_amount"), + refilledAt: bigint("refilled_at", { + mode: "number", + unsigned: true, + }), + ...lifecycleDatesV2, + }, + (table) => ({ + workspaceIdIdx: index("workspace_id_idx").on(table.workspaceId), + uniquePerKey: uniqueIndex("unique_per_key_idx").on(table.keyId), + uniquePerIdentity: uniqueIndex("unique_per_identity_idx").on(table.identityId), + // Index for refill workflow queries + refillLookupIdx: index("refill_lookup_idx").on( + table.refillDay, + table.refillAmount, + table.refilledAt, + ), + }), +); + +export const creditsRelations = relations(credits, ({ one }) => ({ + workspace: one(workspaces, { + fields: [credits.workspaceId], + references: [workspaces.id], + }), + key: one(keys, { + fields: [credits.keyId], + references: [keys.id], + }), + identity: one(identities, { + fields: [credits.identityId], + references: [identities.id], + }), +})); diff --git a/internal/db/src/schema/index.ts b/internal/db/src/schema/index.ts index 94f2eab6d1..eaa22a41ca 100644 --- a/internal/db/src/schema/index.ts +++ b/internal/db/src/schema/index.ts @@ -2,6 +2,7 @@ export * from "./apis"; export * from "./rbac"; export * from "./keyAuth"; export * from "./keys"; +export * from "./credits"; export * from "./vercel_integration"; export * from "./ratelimit"; export * from "./workspaces"; diff --git a/internal/db/src/schema/keys.ts b/internal/db/src/schema/keys.ts index 220d677f3a..a707b59e39 100644 --- a/internal/db/src/schema/keys.ts +++ b/internal/db/src/schema/keys.ts @@ -11,6 +11,7 @@ import { uniqueIndex, varchar, } from "drizzle-orm/mysql-core"; +import { credits } from "./credits"; import { identities, ratelimits } from "./identity"; import { keyAuth } from "./keyAuth"; import { keysPermissions, keysRoles } from "./rbac"; @@ -48,6 +49,7 @@ export const keys = mysqlTable( ...lifecycleDatesMigration, /** + * @deprecated * You can refill uses to keys at a desired interval * * Specify the day on which we should refill. @@ -57,17 +59,34 @@ export const keys = mysqlTable( * - null = we refill on every day */ refillDay: tinyint("refill_day"), + + /** + * @deprecated + * You can refill uses to keys at a desired interval + * + * Specify the amount of uses we should refill. + * - 1 = we refill 1 use + * - 10 = we refill 10 uses + */ refillAmount: int("refill_amount"), + + /** + * @deprecated + * You can refill uses to keys at a desired interval + * + * Specify the last time we refilled the key. + */ lastRefillAt: datetime("last_refill_at", { fsp: 3 }), + /** * sets if key is enabled or disabled */ enabled: boolean("enabled").default(true).notNull(), /** + * @deprecated * You can limit the amount of times a key can be verified before it becomes invalid */ - remaining: int("remaining_requests"), ratelimitAsync: boolean("ratelimit_async"), ratelimitLimit: int("ratelimit_limit"), // max size of the bucket @@ -106,7 +125,6 @@ export const keysRelations = relations(keys, ({ one, many }) => ({ fields: [keys.workspaceId], references: [workspaces.id], }), - forWorkspace: one(workspaces, { fields: [keys.forWorkspaceId], references: [workspaces.id], @@ -119,6 +137,10 @@ export const keysRelations = relations(keys, ({ one, many }) => ({ }), encrypted: one(encryptedKeys), ratelimits: many(ratelimits), + credits: one(credits, { + fields: [keys.id], + references: [credits.keyId], + }), identity: one(identities, { fields: [keys.identityId], references: [identities.id], diff --git a/internal/db/src/types.ts b/internal/db/src/types.ts index 729a356ecd..82bd65cdc2 100644 --- a/internal/db/src/types.ts +++ b/internal/db/src/types.ts @@ -54,3 +54,6 @@ export type InsertAuditLogTarget = InferInsertModel; export type InsertQuotas = InferInsertModel; + +export type Credits = InferSelectModel; +export type InsertCredits = InferInsertModel; diff --git a/internal/id/src/generate.ts b/internal/id/src/generate.ts index 7b6e1e3c2b..f71b8f3833 100644 --- a/internal/id/src/generate.ts +++ b/internal/id/src/generate.ts @@ -4,6 +4,7 @@ const b58 = baseX("123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"); const prefixes = { key: "key", + credit: "cr", policy: "pol", api: "api", request: "req", diff --git a/internal/metrics/src/index.ts b/internal/metrics/src/index.ts index 447946499a..aa2a69faef 100644 --- a/internal/metrics/src/index.ts +++ b/internal/metrics/src/index.ts @@ -86,7 +86,8 @@ export const metricSchema = z.discriminatedUnion("metric", [ }), z.object({ metric: z.literal("metric.usagelimit"), - keyId: z.string(), + keyId: z.string().optional(), + creditId: z.string().optional(), latency: z.number(), }), z.object({