Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 32aee9e

Browse files
pawbanapawbana
authored
feat: add DB queries for ai_gateway_coderd_keys (#25564)
Adds Insert, List and Delete queries for `ai_gateway_coderd_keys ` table.
1 parent 7195be8 commit 32aee9e

8 files changed

Lines changed: 415 additions & 0 deletions

File tree

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1907,6 +1907,13 @@ func (q *querier) CustomRoles(ctx context.Context, arg database.CustomRolesParam
19071907
return q.db.CustomRoles(ctx, arg)
19081908
}
19091909

1910+
func (q *querier) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (database.DeleteAIGatewayKeyRow, error) {
1911+
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAIGatewayKey); err != nil {
1912+
return database.DeleteAIGatewayKeyRow{}, err
1913+
}
1914+
return q.db.DeleteAIGatewayKey(ctx, id)
1915+
}
1916+
19101917
func (q *querier) DeleteAIProviderByID(ctx context.Context, id uuid.UUID) error {
19111918
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAIProvider); err != nil {
19121919
return err
@@ -5463,6 +5470,13 @@ func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.Ins
54635470
return q.db.InsertAIBridgeUserPrompt(ctx, arg)
54645471
}
54655472

5473+
func (q *querier) InsertAIGatewayKey(ctx context.Context, arg database.InsertAIGatewayKeyParams) (database.InsertAIGatewayKeyRow, error) {
5474+
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceAIGatewayKey); err != nil {
5475+
return database.InsertAIGatewayKeyRow{}, err
5476+
}
5477+
return q.db.InsertAIGatewayKey(ctx, arg)
5478+
}
5479+
54665480
func (q *querier) InsertAIProvider(ctx context.Context, arg database.InsertAIProviderParams) (database.AIProvider, error) {
54675481
if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceAIProvider); err != nil {
54685482
return database.AIProvider{}, err
@@ -6238,6 +6252,13 @@ func (q *querier) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context,
62386252
return q.db.ListAIBridgeUserPromptsByInterceptionIDs(ctx, interceptionIDs)
62396253
}
62406254

6255+
func (q *querier) ListAIGatewayKeys(ctx context.Context) ([]database.ListAIGatewayKeysRow, error) {
6256+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIGatewayKey); err != nil {
6257+
return nil, err
6258+
}
6259+
return q.db.ListAIGatewayKeys(ctx)
6260+
}
6261+
62416262
func (q *querier) ListBoundaryLogsBySessionID(ctx context.Context, arg database.ListBoundaryLogsBySessionIDParams) ([]database.BoundaryLog, error) {
62426263
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceBoundaryLog); err != nil {
62436264
return nil, err

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6638,6 +6638,23 @@ func (s *MethodTestSuite) TestAIBridge() {
66386638
dbm.EXPECT().UpdateEncryptedUserAIProviderKey(gomock.Any(), arg).Return(key, nil).AnyTimes()
66396639
check.Args(arg).Asserts(rbac.ResourceAIProvider, policy.ActionUpdate).Returns(key)
66406640
}))
6641+
6642+
s.Run("InsertAIGatewayKey", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
6643+
params := database.InsertAIGatewayKeyParams{}
6644+
row := database.InsertAIGatewayKeyRow{}
6645+
dbm.EXPECT().InsertAIGatewayKey(gomock.Any(), params).Return(row, nil).AnyTimes()
6646+
check.Args(params).Asserts(rbac.ResourceAIGatewayKey, policy.ActionCreate).Returns(row)
6647+
}))
6648+
s.Run("ListAIGatewayKeys", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
6649+
rows := []database.ListAIGatewayKeysRow{}
6650+
dbm.EXPECT().ListAIGatewayKeys(gomock.Any()).Return(rows, nil).AnyTimes()
6651+
check.Args().Asserts(rbac.ResourceAIGatewayKey, policy.ActionRead).Returns(rows)
6652+
}))
6653+
s.Run("DeleteAIGatewayKey", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
6654+
id := uuid.New()
6655+
dbm.EXPECT().DeleteAIGatewayKey(gomock.Any(), id).Return(database.DeleteAIGatewayKeyRow{}, nil).AnyTimes()
6656+
check.Args(id).Asserts(rbac.ResourceAIGatewayKey, policy.ActionDelete).Returns(database.DeleteAIGatewayKeyRow{})
6657+
}))
66416658
}
66426659

66436660
func (s *MethodTestSuite) TestTelemetry() {

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 24 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 45 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier_test.go‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14733,3 +14733,189 @@ func TestSoftDeleteWorkspaceAgentsByWorkspaceID(t *testing.T) {
1473314733
err = db.SoftDeleteWorkspaceAgentsByWorkspaceID(ctx, wsEmpty)
1473414734
require.NoError(t, err)
1473514735
}
14736+
14737+
func TestAIGatewayKeysTableConstraints(t *testing.T) {
14738+
t.Parallel()
14739+
14740+
db, _ := dbtestutil.NewDB(t)
14741+
ctx := testutil.Context(t, testutil.WaitMedium)
14742+
14743+
preExsiting := database.InsertAIGatewayKeyParams{
14744+
ID: uuid.New(),
14745+
Name: "name",
14746+
SecretPrefix: "cgw_test__1",
14747+
HashedSecret: []byte("first-secret"),
14748+
}
14749+
_, err := db.InsertAIGatewayKey(ctx, preExsiting)
14750+
require.NoError(t, err)
14751+
14752+
tests := []struct {
14753+
name string
14754+
params database.InsertAIGatewayKeyParams
14755+
expectUniqueErr database.UniqueConstraint
14756+
expectCheckErr database.CheckConstraint
14757+
}{
14758+
{
14759+
name: "duplicate name",
14760+
params: aiGatewayKeyParams(preExsiting.Name, "cgw_test002"),
14761+
expectUniqueErr: database.UniqueAiGatewayKeysNameIndex,
14762+
},
14763+
{
14764+
name: "duplicate secret prefix",
14765+
params: aiGatewayKeyParams("different-key", preExsiting.SecretPrefix),
14766+
expectUniqueErr: database.UniqueAiGatewayKeysSecretPrefixIndex,
14767+
},
14768+
{
14769+
name: "duplicate hashed secret",
14770+
params: database.InsertAIGatewayKeyParams{ID: uuid.New(), Name: "other-name", SecretPrefix: "cgw_1234567", HashedSecret: preExsiting.HashedSecret},
14771+
expectUniqueErr: database.UniqueAiGatewayKeysHashedSecretIndex,
14772+
},
14773+
{
14774+
name: "empty name",
14775+
params: aiGatewayKeyParams("", "cgw_1234567"),
14776+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14777+
},
14778+
{
14779+
name: "name with trailing dash",
14780+
params: aiGatewayKeyParams("other-name-", "cgw_1234567"),
14781+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14782+
},
14783+
{
14784+
name: "name with consecutive dashes",
14785+
params: aiGatewayKeyParams("other--name", "cgw_1234567"),
14786+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14787+
},
14788+
{
14789+
name: "name with underscore",
14790+
params: aiGatewayKeyParams("other_name", "cgw_1234567"),
14791+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14792+
},
14793+
{
14794+
name: "name with space",
14795+
params: aiGatewayKeyParams("other name", "cgw_1234567"),
14796+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14797+
},
14798+
{
14799+
name: "name with leading dash",
14800+
params: aiGatewayKeyParams("-other-name", "cgw_1234567"),
14801+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14802+
},
14803+
{
14804+
name: "name longer than 64 characters",
14805+
params: aiGatewayKeyParams(strings.Repeat("a", 65), "cgw_1234567"),
14806+
expectCheckErr: database.CheckAiGatewayKeysNameCheck,
14807+
},
14808+
{
14809+
name: "empty secret prefix",
14810+
params: aiGatewayKeyParams("other-name", ""),
14811+
expectCheckErr: database.CheckAiGatewayKeysSecretPrefixCheck,
14812+
},
14813+
{
14814+
name: "invalid secret prefix length",
14815+
params: aiGatewayKeyParams("other-name", "cgw_short"),
14816+
expectCheckErr: database.CheckAiGatewayKeysSecretPrefixCheck,
14817+
},
14818+
{
14819+
name: "empty hashed secret",
14820+
params: database.InsertAIGatewayKeyParams{ID: uuid.New(), Name: "other-name", SecretPrefix: "cgw_1234567"},
14821+
expectCheckErr: database.CheckAiGatewayKeysHashedSecretCheck,
14822+
},
14823+
}
14824+
14825+
for _, tc := range tests {
14826+
t.Run(tc.name, func(t *testing.T) {
14827+
t.Parallel()
14828+
14829+
ctx := testutil.Context(t, testutil.WaitShort)
14830+
14831+
_, err = db.InsertAIGatewayKey(ctx, tc.params)
14832+
require.Error(t, err)
14833+
requireAIGatewayKeysViolation(t, err, tc.expectUniqueErr, tc.expectCheckErr)
14834+
})
14835+
}
14836+
}
14837+
14838+
func TestAIGatewayKeysQueries(t *testing.T) {
14839+
t.Parallel()
14840+
14841+
db, _ := dbtestutil.NewDB(t)
14842+
ctx := testutil.Context(t, testutil.WaitLong)
14843+
14844+
first := aiGatewayKeyParams("first-key", "cgw_first__")
14845+
second := aiGatewayKeyParams("second-key", "cgw_second_")
14846+
second.HashedSecret = []byte("second-secret")
14847+
14848+
firstRow, err := db.InsertAIGatewayKey(ctx, first)
14849+
require.NoError(t, err)
14850+
require.Equal(t, first.ID, firstRow.ID)
14851+
14852+
require.Equal(t, "first-key", firstRow.Name)
14853+
require.Equal(t, first.SecretPrefix, firstRow.SecretPrefix)
14854+
14855+
secondRow, err := db.InsertAIGatewayKey(ctx, second)
14856+
require.NoError(t, err)
14857+
require.Equal(t, second.ID, secondRow.ID)
14858+
14859+
require.Equal(t, "second-key", secondRow.Name)
14860+
require.Equal(t, second.SecretPrefix, secondRow.SecretPrefix)
14861+
14862+
keys, err := db.ListAIGatewayKeys(ctx)
14863+
require.NoError(t, err)
14864+
require.Len(t, keys, 2)
14865+
14866+
requireAIGatewayKeysRow(t, keys[0], first, firstRow.CreatedAt)
14867+
require.False(t, keys[0].LastUsedAt.Valid)
14868+
requireAIGatewayKeysRow(t, keys[1], second, secondRow.CreatedAt)
14869+
require.False(t, keys[1].LastUsedAt.Valid)
14870+
14871+
deleted, err := db.DeleteAIGatewayKey(ctx, first.ID)
14872+
require.NoError(t, err)
14873+
require.Equal(t, first.ID, deleted.ID)
14874+
require.Equal(t, first.Name, deleted.Name)
14875+
require.Equal(t, first.SecretPrefix, deleted.SecretPrefix)
14876+
require.Equal(t, firstRow.CreatedAt, deleted.CreatedAt)
14877+
14878+
_, err = db.DeleteAIGatewayKey(ctx, first.ID)
14879+
require.ErrorIs(t, err, sql.ErrNoRows)
14880+
14881+
keys, err = db.ListAIGatewayKeys(ctx)
14882+
require.NoError(t, err)
14883+
require.Len(t, keys, 1)
14884+
requireAIGatewayKeysRow(t, keys[0], second, secondRow.CreatedAt)
14885+
}
14886+
14887+
func aiGatewayKeyParams(name string, secretPrefix string) database.InsertAIGatewayKeyParams {
14888+
return database.InsertAIGatewayKeyParams{
14889+
ID: uuid.New(),
14890+
Name: name,
14891+
SecretPrefix: secretPrefix,
14892+
HashedSecret: []byte("secret"),
14893+
}
14894+
}
14895+
14896+
func requireAIGatewayKeysRow(t *testing.T, listRow database.ListAIGatewayKeysRow, insertParams database.InsertAIGatewayKeyParams, insertCreatedAt time.Time) {
14897+
t.Helper()
14898+
14899+
require.Equal(t, insertParams.ID, listRow.ID)
14900+
require.Equal(t, insertParams.Name, listRow.Name)
14901+
require.Equal(t, insertParams.SecretPrefix, listRow.SecretPrefix)
14902+
require.Equal(t, insertCreatedAt, listRow.CreatedAt)
14903+
}
14904+
14905+
func requireAIGatewayKeysViolation(
14906+
t *testing.T,
14907+
err error,
14908+
uniqueConstraint database.UniqueConstraint,
14909+
checkConstraint database.CheckConstraint,
14910+
) {
14911+
t.Helper()
14912+
14913+
switch {
14914+
case uniqueConstraint != "":
14915+
require.True(t, database.IsUniqueViolation(err, uniqueConstraint), "expected %q unique violation, got %v", uniqueConstraint, err)
14916+
case checkConstraint != "":
14917+
require.True(t, database.IsCheckViolation(err, checkConstraint), "expected %q check violation, got %v", checkConstraint, err)
14918+
default:
14919+
require.FailNow(t, "test case must expect a constraint error")
14920+
}
14921+
}

0 commit comments

Comments
 (0)