feat: add DB queries for ai_gateway_coderd_keys (#25564)

Adds Insert, List and Delete queries for `ai_gateway_coderd_keys ` table.
2026-09-24 15:04:27 +08:00 · 2026-06-02 13:25:44 +02:00
parent 7195be87b1
commit 32aee9ea4c
8 changed files with 415 additions and 0 deletions
@@ -14733,3 +14733,189 @@ func TestSoftDeleteWorkspaceAgentsByWorkspaceID(t *testing.T) {
 	err = db.SoftDeleteWorkspaceAgentsByWorkspaceID(ctx, wsEmpty)
 	require.NoError(t, err)
 }
+
+func TestAIGatewayKeysTableConstraints(t *testing.T) {
+	t.Parallel()
+
+	db, _ := dbtestutil.NewDB(t)
+	ctx := testutil.Context(t, testutil.WaitMedium)
+
+	preExsiting := database.InsertAIGatewayKeyParams{
+		ID:           uuid.New(),
+		Name:         "name",
+		SecretPrefix: "cgw_test__1",
+		HashedSecret: []byte("first-secret"),
+	}
+	_, err := db.InsertAIGatewayKey(ctx, preExsiting)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name            string
+		params          database.InsertAIGatewayKeyParams
+		expectUniqueErr database.UniqueConstraint
+		expectCheckErr  database.CheckConstraint
+	}{
+		{
+			name:            "duplicate name",
+			params:          aiGatewayKeyParams(preExsiting.Name, "cgw_test002"),
+			expectUniqueErr: database.UniqueAiGatewayKeysNameIndex,
+		},
+		{
+			name:            "duplicate secret prefix",
+			params:          aiGatewayKeyParams("different-key", preExsiting.SecretPrefix),
+			expectUniqueErr: database.UniqueAiGatewayKeysSecretPrefixIndex,
+		},
+		{
+			name:            "duplicate hashed secret",
+			params:          database.InsertAIGatewayKeyParams{ID: uuid.New(), Name: "other-name", SecretPrefix: "cgw_1234567", HashedSecret: preExsiting.HashedSecret},
+			expectUniqueErr: database.UniqueAiGatewayKeysHashedSecretIndex,
+		},
+		{
+			name:           "empty name",
+			params:         aiGatewayKeyParams("", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name with trailing dash",
+			params:         aiGatewayKeyParams("other-name-", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name with consecutive dashes",
+			params:         aiGatewayKeyParams("other--name", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name with underscore",
+			params:         aiGatewayKeyParams("other_name", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name with space",
+			params:         aiGatewayKeyParams("other name", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name with leading dash",
+			params:         aiGatewayKeyParams("-other-name", "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "name longer than 64 characters",
+			params:         aiGatewayKeyParams(strings.Repeat("a", 65), "cgw_1234567"),
+			expectCheckErr: database.CheckAiGatewayKeysNameCheck,
+		},
+		{
+			name:           "empty secret prefix",
+			params:         aiGatewayKeyParams("other-name", ""),
+			expectCheckErr: database.CheckAiGatewayKeysSecretPrefixCheck,
+		},
+		{
+			name:           "invalid secret prefix length",
+			params:         aiGatewayKeyParams("other-name", "cgw_short"),
+			expectCheckErr: database.CheckAiGatewayKeysSecretPrefixCheck,
+		},
+		{
+			name:           "empty hashed secret",
+			params:         database.InsertAIGatewayKeyParams{ID: uuid.New(), Name: "other-name", SecretPrefix: "cgw_1234567"},
+			expectCheckErr: database.CheckAiGatewayKeysHashedSecretCheck,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitShort)
+
+			_, err = db.InsertAIGatewayKey(ctx, tc.params)
+			require.Error(t, err)
+			requireAIGatewayKeysViolation(t, err, tc.expectUniqueErr, tc.expectCheckErr)
+		})
+	}
+}
+
+func TestAIGatewayKeysQueries(t *testing.T) {
+	t.Parallel()
+
+	db, _ := dbtestutil.NewDB(t)
+	ctx := testutil.Context(t, testutil.WaitLong)
+
+	first := aiGatewayKeyParams("first-key", "cgw_first__")
+	second := aiGatewayKeyParams("second-key", "cgw_second_")
+	second.HashedSecret = []byte("second-secret")
+
+	firstRow, err := db.InsertAIGatewayKey(ctx, first)
+	require.NoError(t, err)
+	require.Equal(t, first.ID, firstRow.ID)
+
+	require.Equal(t, "first-key", firstRow.Name)
+	require.Equal(t, first.SecretPrefix, firstRow.SecretPrefix)
+
+	secondRow, err := db.InsertAIGatewayKey(ctx, second)
+	require.NoError(t, err)
+	require.Equal(t, second.ID, secondRow.ID)
+
+	require.Equal(t, "second-key", secondRow.Name)
+	require.Equal(t, second.SecretPrefix, secondRow.SecretPrefix)
+
+	keys, err := db.ListAIGatewayKeys(ctx)
+	require.NoError(t, err)
+	require.Len(t, keys, 2)
+
+	requireAIGatewayKeysRow(t, keys[0], first, firstRow.CreatedAt)
+	require.False(t, keys[0].LastUsedAt.Valid)
+	requireAIGatewayKeysRow(t, keys[1], second, secondRow.CreatedAt)
+	require.False(t, keys[1].LastUsedAt.Valid)
+
+	deleted, err := db.DeleteAIGatewayKey(ctx, first.ID)
+	require.NoError(t, err)
+	require.Equal(t, first.ID, deleted.ID)
+	require.Equal(t, first.Name, deleted.Name)
+	require.Equal(t, first.SecretPrefix, deleted.SecretPrefix)
+	require.Equal(t, firstRow.CreatedAt, deleted.CreatedAt)
+
+	_, err = db.DeleteAIGatewayKey(ctx, first.ID)
+	require.ErrorIs(t, err, sql.ErrNoRows)
+
+	keys, err = db.ListAIGatewayKeys(ctx)
+	require.NoError(t, err)
+	require.Len(t, keys, 1)
+	requireAIGatewayKeysRow(t, keys[0], second, secondRow.CreatedAt)
+}
+
+func aiGatewayKeyParams(name string, secretPrefix string) database.InsertAIGatewayKeyParams {
+	return database.InsertAIGatewayKeyParams{
+		ID:           uuid.New(),
+		Name:         name,
+		SecretPrefix: secretPrefix,
+		HashedSecret: []byte("secret"),
+	}
+}
+
+func requireAIGatewayKeysRow(t *testing.T, listRow database.ListAIGatewayKeysRow, insertParams database.InsertAIGatewayKeyParams, insertCreatedAt time.Time) {
+	t.Helper()
+
+	require.Equal(t, insertParams.ID, listRow.ID)
+	require.Equal(t, insertParams.Name, listRow.Name)
+	require.Equal(t, insertParams.SecretPrefix, listRow.SecretPrefix)
+	require.Equal(t, insertCreatedAt, listRow.CreatedAt)
+}
+
+func requireAIGatewayKeysViolation(
+	t *testing.T,
+	err error,
+	uniqueConstraint database.UniqueConstraint,
+	checkConstraint database.CheckConstraint,
+) {
+	t.Helper()
+
+	switch {
+	case uniqueConstraint != "":
+		require.True(t, database.IsUniqueViolation(err, uniqueConstraint), "expected %q unique violation, got %v", uniqueConstraint, err)
+	case checkConstraint != "":
+		require.True(t, database.IsCheckViolation(err, checkConstraint), "expected %q check violation, got %v", checkConstraint, err)
+	default:
+		require.FailNow(t, "test case must expect a constraint error")
+	}
+}