diff --git a/go/apps/api/routes/v2_apis_list_keys/200_test.go b/go/apps/api/routes/v2_apis_list_keys/200_test.go index 40d6a5587a..53cdc8548c 100644 --- a/go/apps/api/routes/v2_apis_list_keys/200_test.go +++ b/go/apps/api/routes/v2_apis_list_keys/200_test.go @@ -3,7 +3,6 @@ package handler_test import ( "context" "database/sql" - "encoding/json" "fmt" "net/http" "testing" @@ -15,9 +14,8 @@ import ( "github.com/unkeyed/unkey/go/pkg/hash" "github.com/unkeyed/unkey/go/pkg/ptr" "github.com/unkeyed/unkey/go/pkg/testutil" + "github.com/unkeyed/unkey/go/pkg/testutil/seed" "github.com/unkeyed/unkey/go/pkg/uid" - - vaultv1 "github.com/unkeyed/unkey/go/gen/proto/vault/v1" ) func TestSuccess(t *testing.T) { @@ -70,127 +68,75 @@ func TestSuccess(t *testing.T) { require.NoError(t, err) // Create test identities - identity1ID := uid.New("identity") identity1ExternalID := "test_user_1" - err = db.Query.InsertIdentity(ctx, h.DB.RW(), db.InsertIdentityParams{ - ID: identity1ID, - ExternalID: identity1ExternalID, + identity1 := h.CreateIdentity(seed.CreateIdentityRequest{ WorkspaceID: workspace.ID, - Environment: "", - CreatedAt: time.Now().UnixMilli(), + ExternalID: identity1ExternalID, Meta: []byte(`{"role": "admin"}`), }) - require.NoError(t, err) - identity2ID := uid.New("identity") identity2ExternalID := "test_user_2" - err = db.Query.InsertIdentity(ctx, h.DB.RW(), db.InsertIdentityParams{ - ID: identity2ID, - ExternalID: identity2ExternalID, + identity2 := h.CreateIdentity(seed.CreateIdentityRequest{ WorkspaceID: workspace.ID, - Environment: "", - CreatedAt: time.Now().UnixMilli(), + ExternalID: identity2ExternalID, Meta: []byte(`{"role": "user"}`), }) - require.NoError(t, err) - encryptedKeysMap := make(map[string]struct{}) // Create test keys with various configurations - testKeys := []struct { - id string - start string - name string - identityID *string - meta map[string]interface{} - expires *time.Time - enabled bool - }{ - { - id: uid.New("key"), - start: "test_key1_", - name: "Test Key 1", - identityID: &identity1ID, - meta: map[string]interface{}{"env": "production", "team": "backend"}, - enabled: true, - }, - { - id: uid.New("key"), - start: "test_key2_", - name: "Test Key 2", - identityID: &identity1ID, - meta: map[string]interface{}{"env": "staging"}, - enabled: true, - }, - { - id: uid.New("key"), - start: "test_key3_", - name: "Test Key 3", - identityID: &identity2ID, - meta: map[string]interface{}{"env": "development"}, - enabled: true, - }, - { - id: uid.New("key"), - start: "test_key4_", - name: "Test Key 4 (No Identity)", - enabled: true, - }, - { - id: uid.New("key"), - start: "test_key5_", - name: "Test Key 5 (Disabled)", - enabled: false, - }, - } - - for _, keyData := range testKeys { - metaBytes := []byte("{}") - if keyData.meta != nil { - metaBytes, _ = json.Marshal(keyData.meta) - } - - key := keyData.start + uid.New("") + // Track encrypted keys for verification + encryptedKeys := make(map[string]string) // keyID -> plaintext - insertParams := db.InsertKeyParams{ - ID: keyData.id, - KeySpaceID: keySpaceID, - Hash: hash.Sha256(key), - Start: keyData.start, - WorkspaceID: workspace.ID, - ForWorkspaceID: sql.NullString{Valid: false}, - Name: sql.NullString{Valid: true, String: keyData.name}, - Meta: sql.NullString{Valid: true, String: string(metaBytes)}, - Expires: sql.NullTime{Valid: false}, - CreatedAtM: time.Now().UnixMilli(), - Enabled: keyData.enabled, - RemainingRequests: sql.NullInt32{Valid: false}, - } + // Key 1: identity1, production metadata + key1 := h.CreateKey(seed.CreateKeyRequest{ + WorkspaceID: workspace.ID, + KeySpaceID: keySpaceID, + Name: ptr.P("Test Key 1"), + IdentityID: ptr.P(identity1.ID), + Meta: ptr.P(`{"env": "production", "team": "backend"}`), + Recoverable: true, + }) + encryptedKeys[key1.KeyID] = key1.Key - if keyData.identityID != nil { - insertParams.IdentityID = sql.NullString{Valid: true, String: *keyData.identityID} - } else { - insertParams.IdentityID = sql.NullString{Valid: false} - } + // Key 2: identity1, staging metadata + key2 := h.CreateKey(seed.CreateKeyRequest{ + WorkspaceID: workspace.ID, + KeySpaceID: keySpaceID, + Name: ptr.P("Test Key 2"), + IdentityID: ptr.P(identity1.ID), + Meta: ptr.P(`{"env": "staging"}`), + Recoverable: true, + }) + encryptedKeys[key2.KeyID] = key2.Key - err := db.Query.InsertKey(ctx, h.DB.RW(), insertParams) - require.NoError(t, err) + // Key 3: identity2, development metadata + key3 := h.CreateKey(seed.CreateKeyRequest{ + WorkspaceID: workspace.ID, + KeySpaceID: keySpaceID, + Name: ptr.P("Test Key 3"), + IdentityID: ptr.P(identity2.ID), + Meta: ptr.P(`{"env": "development"}`), + Recoverable: true, + }) + encryptedKeys[key3.KeyID] = key3.Key - encryption, err := h.Vault.Encrypt(ctx, &vaultv1.EncryptRequest{ - Keyring: h.Resources().UserWorkspace.ID, - Data: key, - }) - require.NoError(t, err) + // Key 4: no identity + key4 := h.CreateKey(seed.CreateKeyRequest{ + WorkspaceID: workspace.ID, + KeySpaceID: keySpaceID, + Name: ptr.P("Test Key 4 (No Identity)"), + Recoverable: true, + }) + encryptedKeys[key4.KeyID] = key4.Key - err = db.Query.InsertKeyEncryption(ctx, h.DB.RW(), db.InsertKeyEncryptionParams{ - WorkspaceID: h.Resources().UserWorkspace.ID, - KeyID: keyData.id, - CreatedAt: time.Now().UnixMilli(), - Encrypted: encryption.GetEncrypted(), - EncryptionKeyID: encryption.GetKeyId(), - }) - require.NoError(t, err) - encryptedKeysMap[keyData.id] = struct{}{} - } + // Key 5: disabled + key5 := h.CreateKey(seed.CreateKeyRequest{ + WorkspaceID: workspace.ID, + KeySpaceID: keySpaceID, + Name: ptr.P("Test Key 5 (Disabled)"), + Disabled: true, + Recoverable: true, + }) + encryptedKeys[key5.KeyID] = key5.Key // Set up request headers headers := http.Header{ @@ -604,12 +550,13 @@ func TestSuccess(t *testing.T) { require.NotNil(t, res.Body.Data) for _, key := range res.Body.Data { - _, exists := encryptedKeysMap[key.KeyId] + expectedPlaintext, exists := encryptedKeys[key.KeyId] if !exists { continue } require.NotEmpty(t, key.Plaintext, "Key should be decrypted and have plaintext") + require.Equal(t, expectedPlaintext, key.Plaintext, "Key should be decrypted and have correct plaintext") } }) } diff --git a/go/apps/api/routes/v2_apis_list_keys/handler.go b/go/apps/api/routes/v2_apis_list_keys/handler.go index 350988757e..71275b88a7 100644 --- a/go/apps/api/routes/v2_apis_list_keys/handler.go +++ b/go/apps/api/routes/v2_apis_list_keys/handler.go @@ -2,6 +2,7 @@ package handler import ( "context" + "database/sql" "net/http" "github.com/oapi-codegen/nullable" @@ -161,9 +162,35 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { limit := ptr.SafeDeref(req.Limit, 100) cursor := ptr.SafeDeref(req.Cursor, "") - var identityFilter string + // Resolve identity ID if external_id filter is provided + var identityID sql.NullString if req.ExternalId != nil && *req.ExternalId != "" { - identityFilter = *req.ExternalId + identity, identityErr := db.Query.FindIdentityByExternalID(ctx, h.DB.RO(), db.FindIdentityByExternalIDParams{ + WorkspaceID: auth.AuthorizedWorkspaceID, + ExternalID: *req.ExternalId, + Deleted: false, + }) + if identityErr != nil { + if db.IsNotFound(identityErr) { + // Identity doesn't exist, return empty result set + return s.JSON(http.StatusOK, Response{ + Meta: openapi.Meta{ + RequestId: s.RequestID(), + }, + Data: []openapi.KeyResponseData{}, + Pagination: &openapi.Pagination{ + Cursor: nil, + HasMore: false, + }, + }) + } + return fault.Wrap(identityErr, + fault.Code(codes.App.Internal.ServiceUnavailable.URN()), + fault.Internal("database error"), + fault.Public("Failed to retrieve identity."), + ) + } + identityID = sql.NullString{String: identity.ID, Valid: true} } // Query keys by key_auth_id instead of api_id @@ -173,7 +200,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { db.ListLiveKeysByKeySpaceIDParams{ KeySpaceID: api.KeyAuthID.String, IDCursor: cursor, - Identity: identityFilter, + IdentityID: identityID, Limit: int32(limit + 1), // nolint:gosec }, ) 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 42b8db51cf..a3e9177527 100644 --- a/go/apps/api/routes/v2_identities_get_identity/handler.go +++ b/go/apps/api/routes/v2_identities_get_identity/handler.go @@ -50,28 +50,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.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 ratelimits, err := db.UnmarshalNullableJSONTo[[]db.RatelimitInfo](identity.Ratelimits) if err != nil { diff --git a/go/apps/api/routes/v2_identities_update_identity/404_test.go b/go/apps/api/routes/v2_identities_update_identity/404_test.go index 4e6d39c562..2fb61dddf1 100644 --- a/go/apps/api/routes/v2_identities_update_identity/404_test.go +++ b/go/apps/api/routes/v2_identities_update_identity/404_test.go @@ -41,7 +41,7 @@ func TestNotFound(t *testing.T) { res := testutil.CallRoute[handler.Request, openapi.NotFoundErrorResponse](h, route, headers, req) require.Equal(t, http.StatusNotFound, res.Status, "expected 404, got: %d", res.Status) require.Equal(t, "https://unkey.com/docs/errors/unkey/data/identity_not_found", res.Body.Error.Type) - require.Equal(t, "Identity not found in this workspace", res.Body.Error.Detail) + require.Equal(t, "This identity does not exist.", res.Body.Error.Detail) require.Equal(t, http.StatusNotFound, res.Body.Error.Status) require.Equal(t, "Not Found", res.Body.Error.Title) require.NotEmpty(t, res.Body.Meta.RequestId) 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 6130fe09d8..ec254884d6 100644 --- a/go/apps/api/routes/v2_identities_update_identity/handler.go +++ b/go/apps/api/routes/v2_identities_update_identity/handler.go @@ -112,29 +112,25 @@ 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("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("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 +138,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/pkg/db/identity_find.sql_generated.go b/go/pkg/db/identity_find.sql_generated.go new file mode 100644 index 0000000000..95bbfb6917 --- /dev/null +++ b/go/pkg/db/identity_find.sql_generated.go @@ -0,0 +1,118 @@ +// 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, + 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 +JOIN ( + SELECT id1.id FROM identities id1 + WHERE id1.id = ? + AND id1.workspace_id = ? + AND id1.deleted = ? + UNION ALL + SELECT id2.id FROM identities id2 + WHERE id2.workspace_id = ? + AND id2.external_id = ? + AND id2.deleted = ? +) AS identity_lookup ON i.id = identity_lookup.id +LIMIT 1 +` + +type FindIdentityParams struct { + Identity string `db:"identity"` + WorkspaceID string `db:"workspace_id"` + 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"` + 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, +// 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 +// JOIN ( +// SELECT id1.id FROM identities id1 +// WHERE id1.id = ? +// AND id1.workspace_id = ? +// AND id1.deleted = ? +// UNION ALL +// SELECT id2.id FROM identities id2 +// WHERE id2.workspace_id = ? +// AND id2.external_id = ? +// AND id2.deleted = ? +// ) AS identity_lookup ON i.id = identity_lookup.id +// LIMIT 1 +func (q *Queries) FindIdentity(ctx context.Context, db DBTX, arg FindIdentityParams) (FindIdentityRow, error) { + row := db.QueryRowContext(ctx, findIdentity, + arg.Identity, + arg.WorkspaceID, + 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.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/key_list_live_by_key_space_id.sql_generated.go b/go/pkg/db/key_list_live_by_key_space_id.sql_generated.go index d531f85134..4b4522f825 100644 --- a/go/pkg/db/key_list_live_by_key_space_id.sql_generated.go +++ b/go/pkg/db/key_list_live_by_key_space_id.sql_generated.go @@ -90,27 +90,23 @@ SELECT k.id, k.key_auth_id, k.hash, k.start, k.workspace_id, k.for_workspace_id, JSON_ARRAY() ) AS ratelimits FROM ` + "`" + `keys` + "`" + ` k - JOIN key_auth ka ON ka.id = k.key_auth_id - JOIN workspaces ws ON ws.id = k.workspace_id + STRAIGHT_JOIN key_auth ka ON ka.id = k.key_auth_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 WHERE k.key_auth_id = ? AND k.id >= ? - AND ( - ? = '' OR (i.external_id = ? OR i.id = ?) - ) + AND (? IS NULL OR k.identity_id = ?) AND k.deleted_at_m IS NULL AND ka.deleted_at_m IS NULL - AND ws.deleted_at_m IS NULL ORDER BY k.id ASC LIMIT ? ` type ListLiveKeysByKeySpaceIDParams struct { - KeySpaceID string `db:"key_space_id"` - IDCursor string `db:"id_cursor"` - Identity string `db:"identity"` - Limit int32 `db:"limit"` + KeySpaceID string `db:"key_space_id"` + IDCursor string `db:"id_cursor"` + IdentityID sql.NullString `db:"identity_id"` + Limit int32 `db:"limit"` } type ListLiveKeysByKeySpaceIDRow struct { @@ -229,27 +225,22 @@ type ListLiveKeysByKeySpaceIDRow struct { // JSON_ARRAY() // ) AS ratelimits // FROM `keys` k -// JOIN key_auth ka ON ka.id = k.key_auth_id -// JOIN workspaces ws ON ws.id = k.workspace_id +// STRAIGHT_JOIN key_auth ka ON ka.id = k.key_auth_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 // WHERE k.key_auth_id = ? // AND k.id >= ? -// AND ( -// ? = '' OR (i.external_id = ? OR i.id = ?) -// ) +// AND (? IS NULL OR k.identity_id = ?) // AND k.deleted_at_m IS NULL // AND ka.deleted_at_m IS NULL -// AND ws.deleted_at_m IS NULL // ORDER BY k.id ASC // LIMIT ? func (q *Queries) ListLiveKeysByKeySpaceID(ctx context.Context, db DBTX, arg ListLiveKeysByKeySpaceIDParams) ([]ListLiveKeysByKeySpaceIDRow, error) { rows, err := db.QueryContext(ctx, listLiveKeysByKeySpaceID, arg.KeySpaceID, arg.IDCursor, - arg.Identity, - arg.Identity, - arg.Identity, + arg.IdentityID, + arg.IdentityID, arg.Limit, ) if err != nil { diff --git a/go/pkg/db/querier_generated.go b/go/pkg/db/querier_generated.go index a19e3a22c8..bbc1d0170b 100644 --- a/go/pkg/db/querier_generated.go +++ b/go/pkg/db/querier_generated.go @@ -278,23 +278,7 @@ 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, @@ -314,33 +298,35 @@ 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 = ? + // JOIN ( + // SELECT id1.id FROM identities id1 + // WHERE id1.id = ? + // AND id1.workspace_id = ? + // AND id1.deleted = ? + // UNION ALL + // SELECT id2.id FROM identities id2 + // WHERE id2.workspace_id = ? + // AND id2.external_id = ? + // AND id2.deleted = ? + // ) AS identity_lookup ON i.id = identity_lookup.id // LIMIT 1 - FindIdentityWithRatelimits(ctx context.Context, db DBTX, arg FindIdentityWithRatelimitsParams) ([]FindIdentityWithRatelimitsRow, error) + 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 @@ -1582,18 +1568,14 @@ type Querier interface { // JSON_ARRAY() // ) AS ratelimits // FROM `keys` k - // JOIN key_auth ka ON ka.id = k.key_auth_id - // JOIN workspaces ws ON ws.id = k.workspace_id + // STRAIGHT_JOIN key_auth ka ON ka.id = k.key_auth_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 // WHERE k.key_auth_id = ? // AND k.id >= ? - // AND ( - // ? = '' OR (i.external_id = ? OR i.id = ?) - // ) + // AND (? IS NULL OR k.identity_id = ?) // AND k.deleted_at_m IS NULL // AND ka.deleted_at_m IS NULL - // AND ws.deleted_at_m IS NULL // ORDER BY k.id ASC // LIMIT ? ListLiveKeysByKeySpaceID(ctx context.Context, db DBTX, arg ListLiveKeysByKeySpaceIDParams) ([]ListLiveKeysByKeySpaceIDRow, error) diff --git a/go/pkg/db/queries/identity_find.sql b/go/pkg/db/queries/identity_find.sql new file mode 100644 index 0000000000..8e06809cfc --- /dev/null +++ b/go/pkg/db/queries/identity_find.sql @@ -0,0 +1,31 @@ +-- name: FindIdentity :one +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 +JOIN ( + SELECT id1.id FROM identities id1 + WHERE id1.id = sqlc.arg(identity) + AND id1.workspace_id = sqlc.arg(workspace_id) + AND id1.deleted = sqlc.arg(deleted) + UNION ALL + SELECT id2.id FROM identities id2 + WHERE id2.workspace_id = sqlc.arg(workspace_id) + AND id2.external_id = sqlc.arg(identity) + AND id2.deleted = sqlc.arg(deleted) +) AS identity_lookup ON i.id = identity_lookup.id +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/key_list_live_by_key_space_id.sql b/go/pkg/db/queries/key_list_live_by_key_space_id.sql index fc1a53a51d..86f48998cd 100644 --- a/go/pkg/db/queries/key_list_live_by_key_space_id.sql +++ b/go/pkg/db/queries/key_list_live_by_key_space_id.sql @@ -78,17 +78,13 @@ SELECT k.*, JSON_ARRAY() ) AS ratelimits FROM `keys` k - JOIN key_auth ka ON ka.id = k.key_auth_id - JOIN workspaces ws ON ws.id = k.workspace_id + STRAIGHT_JOIN key_auth ka ON ka.id = k.key_auth_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 WHERE k.key_auth_id = sqlc.arg(key_space_id) AND k.id >= sqlc.arg(id_cursor) - AND ( - sqlc.arg(identity) = '' OR (i.external_id = sqlc.arg(identity) OR i.id = sqlc.arg(identity)) - ) + AND (sqlc.arg(identity_id) IS NULL OR k.identity_id = sqlc.arg(identity_id)) AND k.deleted_at_m IS NULL AND ka.deleted_at_m IS NULL - AND ws.deleted_at_m IS NULL ORDER BY k.id ASC LIMIT ?; diff --git a/tools/artillery/create-keys.ts b/tools/artillery/create-keys.ts index 520b798236..a17c90904b 100644 --- a/tools/artillery/create-keys.ts +++ b/tools/artillery/create-keys.ts @@ -152,8 +152,8 @@ async function main() { throw new Error("UNKEY_API_ID not set"); } - const totalKeys = parseInt(process.env.KEY_COUNT || "10000", 10); - const concurrency = parseInt(process.env.CONCURRENCY || "50", 10); + const totalKeys = Number.parseInt(process.env.KEY_COUNT || "10000", 10); + const concurrency = Number.parseInt(process.env.CONCURRENCY || "50", 10); console.info(`Creating ${totalKeys} keys with concurrency ${concurrency}...`);