mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +08:00
feat: Add full text search over chat messages (#27126)
Closes CODAGT-721 Closes CODAGT-722 Closes CODAGT-723 Closes CODAGT-724 Closes CODAGT-725 This PR adds the database and API pieces necessary to support full-text chat message search. - Adds required chat schema for full-text search - Adds dbpurge job to populate search_tsv in the background - Adds `search` parameter to GetChats query - Adds `search` filter to `searchquery.Chats` - Wires chat search filter into chats API > Implemented by Coder Agents, reviewed and tested by a human.
This commit is contained in:
@@ -697,7 +697,8 @@ var (
|
||||
rbac.ResourceApiKey.Type: {policy.ActionDelete},
|
||||
rbac.ResourceAibridgeInterception.Type: {policy.ActionDelete},
|
||||
rbac.ResourceWorkspaceBuildOrchestration.Type: {policy.ActionDelete},
|
||||
// Chat auto-archive sets archived=true on inactive chats.
|
||||
// Chat auto-archive sets archived=true on inactive chats and computes
|
||||
// search_tsv tsvector for chat_messages.
|
||||
rbac.ResourceChat.Type: {policy.ActionRead, policy.ActionUpdate},
|
||||
// Purge old boundary logs past the retention period.
|
||||
rbac.ResourceBoundaryLog.Type: {policy.ActionDelete},
|
||||
@@ -1743,6 +1744,13 @@ func (q *querier) AutoArchiveInactiveChats(ctx context.Context, arg database.Aut
|
||||
return q.db.AutoArchiveInactiveChats(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.BackfillChatMessagesSearchTsv(ctx, batchSize)
|
||||
}
|
||||
|
||||
func (q *querier) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
// This is a system-level operation used by the gitsync
|
||||
// background worker to reschedule failed refreshes. Same
|
||||
@@ -1820,6 +1828,13 @@ func (q *querier) CalculateAIBridgeInterceptionsTelemetrySummary(ctx context.Con
|
||||
return q.db.CalculateAIBridgeInterceptionsTelemetrySummary(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil {
|
||||
return false, err
|
||||
}
|
||||
return q.db.ChatSearchQueryIsEmpty(ctx, search)
|
||||
}
|
||||
|
||||
func (q *querier) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
|
||||
empty := database.ClaimPrebuiltWorkspaceRow{}
|
||||
|
||||
|
||||
@@ -1002,6 +1002,14 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().DeleteOldChats(gomock.Any(), database.DeleteOldChatsParams{}).Return(int64(0), nil).AnyTimes()
|
||||
check.Args(database.DeleteOldChatsParams{}).Asserts(rbac.ResourceSystem, policy.ActionDelete)
|
||||
}))
|
||||
s.Run("BackfillChatMessagesSearchTsv", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), int32(100)).Return(int64(0), nil).AnyTimes()
|
||||
check.Args(int32(100)).Asserts(rbac.ResourceChat, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("ChatSearchQueryIsEmpty", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().ChatSearchQueryIsEmpty(gomock.Any(), "!!!").Return(true, nil).AnyTimes()
|
||||
check.Args("!!!").Asserts(rbac.ResourceChat, policy.ActionRead)
|
||||
}))
|
||||
s.Run("GetChatRetentionDays", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(30), nil).AnyTimes()
|
||||
check.Args().Asserts()
|
||||
|
||||
+16
@@ -177,6 +177,14 @@ func (m queryMetricsStore) AutoArchiveInactiveChats(ctx context.Context, arg dat
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.BackfillChatMessagesSearchTsv(ctx, batchSize)
|
||||
m.queryLatencies.WithLabelValues("BackfillChatMessagesSearchTsv").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BackfillChatMessagesSearchTsv").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.BackoffChatDiffStatus(ctx, arg)
|
||||
@@ -257,6 +265,14 @@ func (m queryMetricsStore) CalculateAIBridgeInterceptionsTelemetrySummary(ctx co
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ChatSearchQueryIsEmpty(ctx, search)
|
||||
m.queryLatencies.WithLabelValues("ChatSearchQueryIsEmpty").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ChatSearchQueryIsEmpty").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ClaimPrebuiltWorkspace(ctx, arg)
|
||||
|
||||
Generated
+30
@@ -178,6 +178,21 @@ func (mr *MockStoreMockRecorder) AutoArchiveInactiveChats(ctx, arg any) *gomock.
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AutoArchiveInactiveChats", reflect.TypeOf((*MockStore)(nil).AutoArchiveInactiveChats), ctx, arg)
|
||||
}
|
||||
|
||||
// BackfillChatMessagesSearchTsv mocks base method.
|
||||
func (m *MockStore) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "BackfillChatMessagesSearchTsv", ctx, batchSize)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// BackfillChatMessagesSearchTsv indicates an expected call of BackfillChatMessagesSearchTsv.
|
||||
func (mr *MockStoreMockRecorder) BackfillChatMessagesSearchTsv(ctx, batchSize any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BackfillChatMessagesSearchTsv", reflect.TypeOf((*MockStore)(nil).BackfillChatMessagesSearchTsv), ctx, batchSize)
|
||||
}
|
||||
|
||||
// BackoffChatDiffStatus mocks base method.
|
||||
func (m *MockStore) BackoffChatDiffStatus(ctx context.Context, arg database.BackoffChatDiffStatusParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -322,6 +337,21 @@ func (mr *MockStoreMockRecorder) CalculateAIBridgeInterceptionsTelemetrySummary(
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CalculateAIBridgeInterceptionsTelemetrySummary", reflect.TypeOf((*MockStore)(nil).CalculateAIBridgeInterceptionsTelemetrySummary), ctx, arg)
|
||||
}
|
||||
|
||||
// ChatSearchQueryIsEmpty mocks base method.
|
||||
func (m *MockStore) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ChatSearchQueryIsEmpty", ctx, search)
|
||||
ret0, _ := ret[0].(bool)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ChatSearchQueryIsEmpty indicates an expected call of ChatSearchQueryIsEmpty.
|
||||
func (mr *MockStoreMockRecorder) ChatSearchQueryIsEmpty(ctx, search any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ChatSearchQueryIsEmpty", reflect.TypeOf((*MockStore)(nil).ChatSearchQueryIsEmpty), ctx, search)
|
||||
}
|
||||
|
||||
// ClaimPrebuiltWorkspace mocks base method.
|
||||
func (m *MockStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
@@ -51,6 +51,12 @@ const (
|
||||
// Chat debug run deletions can cascade into steps with large JSONB
|
||||
// payloads, so they use the same conservative batch size.
|
||||
chatDebugRunsBatchSize = 1000
|
||||
// Chat search tsvector backfill is capped at 5 batches of 10k
|
||||
// rows per tick. Benchmarks on a dogfood-class machine (EPYC 9454P)
|
||||
// with containerized Postgres were measured to take ~800ms per batch.
|
||||
// This is considered acceptable but may need dialing in later.
|
||||
chatSearchBackfillBatchSize = 10000
|
||||
chatSearchBackfillMaxBatches = 5
|
||||
)
|
||||
|
||||
type Option func(*instance)
|
||||
@@ -61,6 +67,14 @@ func WithClock(clk quartz.Clock) Option {
|
||||
return func(i *instance) { i.clk = clk }
|
||||
}
|
||||
|
||||
// WithChatSearchBackfillLimits overrides backfill batch size and cap. For tests.
|
||||
func WithChatSearchBackfillLimits(batchSize int32, maxBatches int) Option {
|
||||
return func(i *instance) {
|
||||
i.chatSearchBackfillBatchSize = batchSize
|
||||
i.chatSearchBackfillMaxBatches = maxBatches
|
||||
}
|
||||
}
|
||||
|
||||
// New creates a new periodically purging database instance.
|
||||
// Callers must Close the returned instance.
|
||||
func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, opts ...Option) io.Closer {
|
||||
@@ -87,14 +101,25 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder
|
||||
}, []string{"record_type"})
|
||||
reg.MustRegister(recordsPurged)
|
||||
|
||||
chatSearchRowsBackfilled := prometheus.NewCounter(prometheus.CounterOpts{
|
||||
Namespace: "coderd",
|
||||
Subsystem: "dbpurge",
|
||||
Name: "chat_search_rows_backfilled_total",
|
||||
Help: "Total number of chat message rows whose search_tsv was backfilled.",
|
||||
})
|
||||
reg.MustRegister(chatSearchRowsBackfilled)
|
||||
|
||||
inst := &instance{
|
||||
cancel: cancelFunc,
|
||||
closed: closed,
|
||||
logger: logger,
|
||||
vals: vals,
|
||||
clk: quartz.NewReal(),
|
||||
iterationDuration: iterationDuration,
|
||||
recordsPurged: recordsPurged,
|
||||
cancel: cancelFunc,
|
||||
closed: closed,
|
||||
logger: logger,
|
||||
vals: vals,
|
||||
clk: quartz.NewReal(),
|
||||
iterationDuration: iterationDuration,
|
||||
recordsPurged: recordsPurged,
|
||||
chatSearchRowsBackfilled: chatSearchRowsBackfilled,
|
||||
chatSearchBackfillBatchSize: chatSearchBackfillBatchSize,
|
||||
chatSearchBackfillMaxBatches: chatSearchBackfillMaxBatches,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(inst)
|
||||
@@ -310,6 +335,25 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
}
|
||||
}
|
||||
|
||||
// Backfill search_tsv tsvector on chat_messages in batches. Doing this here because it's
|
||||
// potentially too much for a regular migration, especially on larger deployments:
|
||||
// - Each row with search_tsv = NULL is present in idx_chat_messages_search_tsv_pending.
|
||||
// - Content of chat_messages is not changed after insert.
|
||||
// - Rows that are soft-deleted are no longer part of the index.
|
||||
// NOTE: This should not remain in dbpurge and should be adjusted when the "DBOps" gets
|
||||
// implemented.
|
||||
var backfilledChatSearchRows int64
|
||||
for range i.chatSearchBackfillMaxBatches {
|
||||
n, err := tx.BackfillChatMessagesSearchTsv(ctx, i.chatSearchBackfillBatchSize)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("backfill chat_messages.search_tsv: %w", err)
|
||||
}
|
||||
backfilledChatSearchRows += n
|
||||
if n < int64(i.chatSearchBackfillBatchSize) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
i.logger.Debug(ctx, "purged old database entries",
|
||||
slog.F("workspace_agent_logs", purgedWorkspaceAgentLogs),
|
||||
slog.F("expired_api_keys", expiredAPIKeys),
|
||||
@@ -322,6 +366,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
slog.F("chats", purgedChats),
|
||||
slog.F("chat_files", purgedChatFiles),
|
||||
slog.F("chat_debug_runs", purgedChatDebugRuns),
|
||||
slog.F("chat_search_rows_backfilled", backfilledChatSearchRows),
|
||||
slog.F("duration", i.clk.Since(start)),
|
||||
)
|
||||
|
||||
@@ -338,6 +383,9 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
i.recordsPurged.WithLabelValues("chat_debug_runs").Add(float64(purgedChatDebugRuns))
|
||||
i.recordsPurged.WithLabelValues("chat_files").Add(float64(purgedChatFiles))
|
||||
}
|
||||
if i.chatSearchRowsBackfilled != nil {
|
||||
i.chatSearchRowsBackfilled.Add(float64(backfilledChatSearchRows))
|
||||
}
|
||||
|
||||
// chatConfigErr is returned after the tx, so do not record this
|
||||
// iteration as successful when only the deferred config read failed.
|
||||
@@ -362,13 +410,16 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.
|
||||
}
|
||||
|
||||
type instance struct {
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
logger slog.Logger
|
||||
vals *codersdk.DeploymentValues
|
||||
clk quartz.Clock
|
||||
iterationDuration *prometheus.HistogramVec
|
||||
recordsPurged *prometheus.CounterVec
|
||||
cancel context.CancelFunc
|
||||
closed chan struct{}
|
||||
logger slog.Logger
|
||||
vals *codersdk.DeploymentValues
|
||||
clk quartz.Clock
|
||||
iterationDuration *prometheus.HistogramVec
|
||||
recordsPurged *prometheus.CounterVec
|
||||
chatSearchRowsBackfilled prometheus.Counter
|
||||
chatSearchBackfillBatchSize int32
|
||||
chatSearchBackfillMaxBatches int
|
||||
}
|
||||
|
||||
func (i *instance) Close() error {
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/google/uuid"
|
||||
"github.com/lib/pq"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/goleak"
|
||||
@@ -254,6 +255,7 @@ func TestMetrics(t *testing.T) {
|
||||
mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldChatDebugRuns(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatDebugRunsParams{})).Return(int64(0), nil).MinTimes(1)
|
||||
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
|
||||
DoAndReturn(func(f func(database.Store) error, _ *database.TxOptions) error {
|
||||
@@ -305,6 +307,7 @@ func TestMetrics(t *testing.T) {
|
||||
mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
|
||||
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes()
|
||||
mDB.EXPECT().DeleteOldChats(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatsParams{})).Return(int64(0), nil).MinTimes(1)
|
||||
mDB.EXPECT().DeleteOldChatFiles(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatFilesParams{})).Return(int64(0), nil).MinTimes(1)
|
||||
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
|
||||
@@ -2861,3 +2864,375 @@ func TestDeleteOldChatFiles(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func awaitDoTicks(ctx context.Context, t *testing.T, clk *quartz.Mock, n int) func() {
|
||||
t.Helper()
|
||||
completed := make(chan struct{})
|
||||
advance := make(chan struct{})
|
||||
trapNow := clk.Trap().Now()
|
||||
trapStop := clk.Trap().TickerStop()
|
||||
trapReset := clk.Trap().TickerReset()
|
||||
go func() {
|
||||
defer close(completed)
|
||||
defer trapReset.Close()
|
||||
defer trapStop.Close()
|
||||
defer trapNow.Close()
|
||||
trapNow.MustWait(ctx).MustRelease(ctx)
|
||||
trapReset.MustWait(ctx).MustRelease(ctx)
|
||||
select {
|
||||
case completed <- struct{}{}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
for i := 1; i < n; i++ {
|
||||
select {
|
||||
case <-advance:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
d, w := clk.AdvanceNext()
|
||||
if !assert.Equal(t, 10*time.Minute, d) {
|
||||
return
|
||||
}
|
||||
w.MustWait(ctx)
|
||||
trapStop.MustWait(ctx).MustRelease(ctx)
|
||||
trapReset.MustWait(ctx).MustRelease(ctx)
|
||||
select {
|
||||
case completed <- struct{}{}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
first := true
|
||||
return func() {
|
||||
t.Helper()
|
||||
if !first {
|
||||
testutil.RequireSend(ctx, t, advance, struct{}{})
|
||||
}
|
||||
first = false
|
||||
testutil.TryReceive(ctx, t, completed)
|
||||
}
|
||||
}
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
func TestBackfillChatMessagesSearchTsv(t *testing.T) {
|
||||
now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
type chatSearchDeps struct {
|
||||
user database.User
|
||||
modelConfig database.ChatModelConfig
|
||||
chat database.Chat
|
||||
}
|
||||
setupDeps := func(t *testing.T, db database.Store) chatSearchDeps {
|
||||
t.Helper()
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
|
||||
UserID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
_ = dbgen.ChatProvider(t, db, database.ChatProvider{
|
||||
Provider: "openai",
|
||||
DisplayName: "OpenAI",
|
||||
})
|
||||
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
})
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelConfig.ID,
|
||||
Title: "search-backfill-test-chat",
|
||||
})
|
||||
return chatSearchDeps{user: user, modelConfig: modelConfig, chat: chat}
|
||||
}
|
||||
textContent := func(text string) pqtype.NullRawMessage {
|
||||
return pqtype.NullRawMessage{
|
||||
RawMessage: json.RawMessage(fmt.Sprintf(`[{"type":"text","text":%q}]`, text)),
|
||||
Valid: true,
|
||||
}
|
||||
}
|
||||
createMessage := func(t *testing.T, db database.Store, deps chatSearchDeps, role database.ChatMessageRole, visibility database.ChatMessageVisibility, content pqtype.NullRawMessage) database.ChatMessage {
|
||||
t.Helper()
|
||||
return dbgen.ChatMessage(t, db, database.ChatMessage{
|
||||
ChatID: deps.chat.ID,
|
||||
CreatedBy: uuid.NullUUID{UUID: deps.user.ID, Valid: true},
|
||||
ModelConfigID: uuid.NullUUID{UUID: deps.modelConfig.ID, Valid: true},
|
||||
Role: role,
|
||||
Visibility: visibility,
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
softDelete := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64) {
|
||||
t.Helper()
|
||||
_, err := rawDB.ExecContext(ctx, "UPDATE chat_messages SET deleted = true WHERE id = $1", id)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
// The WHERE clause below must match the predicate of idx_chat_messages_search_tsv_pending.
|
||||
countPending := func(ctx context.Context, t *testing.T, rawDB *sql.DB) int {
|
||||
t.Helper()
|
||||
var count int
|
||||
err := rawDB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*) FROM chat_messages
|
||||
WHERE search_tsv IS NULL
|
||||
AND deleted = false
|
||||
AND visibility IN ('user', 'both')
|
||||
AND role IN ('user', 'assistant')`).Scan(&count)
|
||||
require.NoError(t, err)
|
||||
return count
|
||||
}
|
||||
searchTsv := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64) (isNull bool, text string) {
|
||||
t.Helper()
|
||||
err := rawDB.QueryRowContext(ctx,
|
||||
"SELECT search_tsv IS NULL, COALESCE(search_tsv::text, '') FROM chat_messages WHERE id = $1", id).
|
||||
Scan(&isNull, &text)
|
||||
require.NoError(t, err)
|
||||
return isNull, text
|
||||
}
|
||||
requireBackfilled := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, msg string) {
|
||||
t.Helper()
|
||||
isNull, _ := searchTsv(ctx, t, rawDB, id)
|
||||
require.False(t, isNull, msg)
|
||||
}
|
||||
// Asserts the row's tsvector matches expectedText, not just non-NULL.
|
||||
requireTsvFor := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, expectedText string) {
|
||||
t.Helper()
|
||||
var matches bool
|
||||
err := rawDB.QueryRowContext(ctx,
|
||||
"SELECT search_tsv = to_tsvector('simple', $2::text) FROM chat_messages WHERE id = $1", id, expectedText).
|
||||
Scan(&matches)
|
||||
require.NoError(t, err)
|
||||
require.True(t, matches, "search_tsv should contain the lexemes of %q", expectedText)
|
||||
}
|
||||
requireNotBackfilled := func(ctx context.Context, t *testing.T, rawDB *sql.DB, id int64, msg string) {
|
||||
t.Helper()
|
||||
isNull, _ := searchTsv(ctx, t, rawDB, id)
|
||||
require.True(t, isNull, msg)
|
||||
}
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("DrainConverges", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
eligibleBoth := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("hello world"))
|
||||
eligibleUserVis := createMessage(t, db, deps, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, textContent("assistant reply"))
|
||||
eligibleNoText := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[]`), Valid: true})
|
||||
toolMsg := createMessage(t, db, deps, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output"))
|
||||
modelOnlyMsg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, textContent("model only"))
|
||||
deletedMsg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deleted message"))
|
||||
softDelete(ctx, t, rawDB, deletedMsg.ID)
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
require.Zero(t, countPending(ctx, t, rawDB), "queue should be drained")
|
||||
requireTsvFor(ctx, t, rawDB, eligibleBoth.ID, "hello world")
|
||||
requireTsvFor(ctx, t, rawDB, eligibleUserVis.ID, "assistant reply")
|
||||
requireBackfilled(ctx, t, rawDB, eligibleNoText.ID, "eligible message with no text should be backfilled (sentinel)")
|
||||
requireNotBackfilled(ctx, t, rawDB, toolMsg.ID, "tool message should never be backfilled")
|
||||
requireNotBackfilled(ctx, t, rawDB, modelOnlyMsg.ID, "model-only message should never be backfilled")
|
||||
requireNotBackfilled(ctx, t, rawDB, deletedMsg.ID, "deleted message should never be backfilled")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("BackfillsNewestFirst", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
var ids []int64
|
||||
for i := range 5 {
|
||||
msg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
|
||||
ids = append(ids, msg.ID)
|
||||
}
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(),
|
||||
dbpurge.WithClock(clk), dbpurge.WithChatSearchBackfillLimits(2, 1))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
slices.Sort(ids)
|
||||
requireBackfilled(ctx, t, rawDB, ids[4], "newest message should be backfilled first")
|
||||
requireBackfilled(ctx, t, rawDB, ids[3], "second-newest message should be backfilled first")
|
||||
for _, id := range ids[:3] {
|
||||
requireNotBackfilled(ctx, t, rawDB, id, "older messages should remain pending after one batch")
|
||||
}
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("NoTextSentinel", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
emptyArr := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[]`), Valid: true})
|
||||
noTextParts := createMessage(t, db, deps, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"tool_call","id":"x"}]`), Valid: true})
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
for _, id := range []int64{emptyArr.ID, noTextParts.ID} {
|
||||
isNull, text := searchTsv(ctx, t, rawDB, id)
|
||||
require.False(t, isNull, "no-text row should get the empty-tsvector sentinel, not stay NULL")
|
||||
require.Empty(t, text, "no-text row should have an empty tsvector")
|
||||
}
|
||||
require.Zero(t, countPending(ctx, t, rawDB), "sentinel rows should not reappear as pending")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("PerTickBound", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
for i := range 6 {
|
||||
createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
|
||||
}
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 2)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(),
|
||||
dbpurge.WithClock(clk), dbpurge.WithChatSearchBackfillLimits(2, 2))
|
||||
defer closer.Close()
|
||||
|
||||
tick()
|
||||
require.Equal(t, 2, countPending(ctx, t, rawDB), "one tick backfills at most maxBatches*batchSize rows")
|
||||
|
||||
tick()
|
||||
require.Zero(t, countPending(ctx, t, rawDB), "next tick continues draining")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("SkipsDeletedRows", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
msg := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("soft deleted before backfill"))
|
||||
softDelete(ctx, t, rawDB, msg.ID)
|
||||
require.Zero(t, countPending(ctx, t, rawDB), "deleted rows should not appear as pending")
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
requireNotBackfilled(ctx, t, rawDB, msg.ID, "deleted row should never be backfilled")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("BackfillsNewMessagesAfterDrain", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
|
||||
initial := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("initial message"))
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 2)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
|
||||
tick()
|
||||
requireBackfilled(ctx, t, rawDB, initial.ID, "initial message should be backfilled")
|
||||
require.Zero(t, countPending(ctx, t, rawDB))
|
||||
|
||||
fresh := createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("post drain message"))
|
||||
tick()
|
||||
requireBackfilled(ctx, t, rawDB, fresh.ID, "message inserted after drain should be backfilled on the next tick")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("SteadyStateNoop", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
_ = setupDeps(t, db)
|
||||
reg := prometheus.NewRegistry()
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, reg, dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
require.Zero(t, countPending(ctx, t, rawDB))
|
||||
backfilled := promhelp.CounterValue(t, reg, "coderd_dbpurge_chat_search_rows_backfilled_total", nil)
|
||||
require.Zero(t, backfilled, "empty queue should backfill zero rows")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("MetricsCountsBackfilledRows", func(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
clk := quartz.NewMock(t)
|
||||
clk.Set(now).MustWait(ctx)
|
||||
db, _, _ := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure())
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
deps := setupDeps(t, db)
|
||||
reg := prometheus.NewRegistry()
|
||||
|
||||
for i := range 3 {
|
||||
createMessage(t, db, deps, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent(fmt.Sprintf("message %d", i)))
|
||||
}
|
||||
createMessage(t, db, deps, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output"))
|
||||
|
||||
tick := awaitDoTicks(ctx, t, clk, 1)
|
||||
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, reg, dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
tick()
|
||||
|
||||
backfilled := promhelp.CounterValue(t, reg, "coderd_dbpurge_chat_search_rows_backfilled_total", nil)
|
||||
require.Equal(t, 3, backfilled, "counter should count exactly the eligible backfilled rows")
|
||||
})
|
||||
|
||||
//nolint:paralleltest // It uses LockIDDBPurge.
|
||||
t.Run("SkippedWhenLockHeld", func(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
|
||||
clk := quartz.NewMock(t)
|
||||
ctrl := gomock.NewController(t)
|
||||
mDB := dbmock.NewMockStore(ctrl)
|
||||
mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(0), nil).AnyTimes()
|
||||
mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays).
|
||||
Return(int32(0), nil).AnyTimes()
|
||||
mDB.EXPECT().TryAcquireLock(gomock.Any(), int64(database.LockIDDBPurge)).Return(false, nil).AnyTimes()
|
||||
mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Times(0)
|
||||
mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).
|
||||
DoAndReturn(func(f func(database.Store) error, _ *database.TxOptions) error {
|
||||
return f(mDB)
|
||||
}).MinTimes(1)
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
done := awaitDoTick(ctx, t, clk)
|
||||
closer := dbpurge.New(ctx, logger, mDB, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
|
||||
defer closer.Close()
|
||||
testutil.TryReceive(ctx, t, done)
|
||||
})
|
||||
}
|
||||
|
||||
Generated
+37
-3
@@ -783,6 +783,18 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE FUNCTION chat_message_search_text(content jsonb) RETURNS text
|
||||
LANGUAGE sql IMMUTABLE PARALLEL SAFE
|
||||
AS $$
|
||||
SELECT CASE WHEN jsonb_typeof(content) = 'array' THEN (
|
||||
SELECT string_agg(part->>'text', ' ' ORDER BY ordinality)
|
||||
FROM jsonb_array_elements(content) WITH ORDINALITY AS t(part, ordinality)
|
||||
WHERE part->>'type' = 'text'
|
||||
) END
|
||||
$$;
|
||||
|
||||
COMMENT ON FUNCTION chat_message_search_text(content jsonb) IS 'Extracts searchable content from chat_messages. Returns NULL for scalar JSON strings (content_version=0). Immutable as it is used in indexes.';
|
||||
|
||||
CREATE FUNCTION check_workspace_agent_name_unique() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
@@ -1365,6 +1377,7 @@ CREATE FUNCTION set_chat_message_revision_before() RETURNS trigger
|
||||
AS $$
|
||||
DECLARE
|
||||
chat_snapshot_version bigint;
|
||||
cmp chat_messages;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
@@ -1379,7 +1392,9 @@ BEGIN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF OLD IS NOT DISTINCT FROM NEW THEN
|
||||
cmp := NEW;
|
||||
cmp.search_tsv := OLD.search_tsv;
|
||||
IF OLD IS NOT DISTINCT FROM cmp THEN
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
END IF;
|
||||
@@ -1396,6 +1411,8 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
COMMENT ON FUNCTION set_chat_message_revision_before() IS 'Component of chatd. Updates chat_snapshot_version when any fields of chat_messages change. Excludes changes to search_tsv as it is not relevant to chatd''s processing loop.';
|
||||
|
||||
CREATE FUNCTION sync_chat_retry_state() RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
@@ -1446,7 +1463,7 @@ BEGIN
|
||||
SELECT DISTINCT n.chat_id
|
||||
FROM chat_message_history_new_rows n
|
||||
JOIN chat_message_history_old_rows o ON o.id = n.id
|
||||
WHERE o IS DISTINCT FROM n
|
||||
WHERE (to_jsonb(o) - 'search_tsv') IS DISTINCT FROM (to_jsonb(n) - 'search_tsv')
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
@@ -1457,6 +1474,8 @@ BEGIN
|
||||
END;
|
||||
$$;
|
||||
|
||||
COMMENT ON FUNCTION update_chat_history_after_message_update() IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv.';
|
||||
|
||||
CREATE TABLE ai_gateway_keys (
|
||||
id uuid NOT NULL,
|
||||
created_at timestamp with time zone NOT NULL,
|
||||
@@ -1946,11 +1965,14 @@ CREATE TABLE chat_messages (
|
||||
provider_response_id text,
|
||||
api_key_id text,
|
||||
revision bigint NOT NULL,
|
||||
reasoning_effort chat_reasoning_effort
|
||||
reasoning_effort chat_reasoning_effort,
|
||||
search_tsv tsvector
|
||||
);
|
||||
|
||||
COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.';
|
||||
|
||||
COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.';
|
||||
|
||||
CREATE SEQUENCE chat_messages_id_seq
|
||||
START WITH 1
|
||||
INCREMENT BY 1
|
||||
@@ -4717,6 +4739,8 @@ CREATE UNIQUE INDEX idx_chat_debug_steps_run_step ON chat_debug_steps USING btre
|
||||
|
||||
CREATE INDEX idx_chat_debug_steps_stale ON chat_debug_steps USING btree (updated_at) WHERE (finished_at IS NULL);
|
||||
|
||||
CREATE INDEX idx_chat_diff_statuses_pr_title_fts ON chat_diff_statuses USING gin (to_tsvector('simple'::regconfig, pull_request_title));
|
||||
|
||||
CREATE INDEX idx_chat_diff_statuses_stale_at ON chat_diff_statuses USING btree (stale_at);
|
||||
|
||||
CREATE INDEX idx_chat_diff_statuses_url_lower ON chat_diff_statuses USING btree (lower(url)) WHERE ((url IS NOT NULL) AND (url <> ''::text));
|
||||
@@ -4737,6 +4761,12 @@ CREATE INDEX idx_chat_messages_created_at ON chat_messages USING btree (created_
|
||||
|
||||
CREATE INDEX idx_chat_messages_owner_spend ON chat_messages USING btree (chat_id, created_at) WHERE (total_cost_micros IS NOT NULL);
|
||||
|
||||
CREATE INDEX idx_chat_messages_search_tsv ON chat_messages USING gin (search_tsv) WHERE ((search_tsv IS NOT NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
|
||||
|
||||
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for populating search_tsv in the background. Only defined over ''searchable'' rows of chat_messages where search_tsv is NULL.';
|
||||
|
||||
CREATE INDEX idx_chat_messages_search_tsv_pending ON chat_messages USING btree (id DESC) WHERE ((search_tsv IS NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
|
||||
|
||||
CREATE INDEX idx_chat_messages_user_prompts ON chat_messages USING btree (chat_id, id DESC) WHERE ((deleted = false) AND (role = 'user'::chat_message_role) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])));
|
||||
|
||||
CREATE INDEX idx_chat_model_configs_ai_provider_id ON chat_model_configs USING btree (ai_provider_id);
|
||||
@@ -4763,6 +4793,10 @@ CREATE INDEX idx_chats_parent_chat_id ON chats USING btree (parent_chat_id);
|
||||
|
||||
CREATE INDEX idx_chats_root_chat_id ON chats USING btree (root_chat_id);
|
||||
|
||||
CREATE INDEX idx_chats_title_fts ON chats USING gin (to_tsvector('simple'::regconfig, title));
|
||||
|
||||
COMMENT ON INDEX idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
|
||||
|
||||
CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (status, updated_at, id) WHERE (archived = false);
|
||||
|
||||
CREATE INDEX idx_chats_workspace ON chats USING btree (workspace_id);
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
-- Restore the original trigger bodies from 000519.
|
||||
CREATE OR REPLACE FUNCTION set_chat_message_revision_before()
|
||||
RETURNS trigger AS $$
|
||||
DECLARE
|
||||
chat_snapshot_version bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'UPDATE' THEN
|
||||
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
|
||||
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
|
||||
END IF;
|
||||
|
||||
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF OLD IS NOT DISTINCT FROM NEW THEN
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
SELECT snapshot_version INTO chat_snapshot_version
|
||||
FROM chats WHERE id = NEW.chat_id;
|
||||
|
||||
IF chat_snapshot_version IS NULL THEN
|
||||
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
NEW.revision = chat_snapshot_version;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
CREATE OR REPLACE FUNCTION update_chat_history_after_message_update()
|
||||
RETURNS trigger AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT n.chat_id
|
||||
FROM chat_message_history_new_rows n
|
||||
JOIN chat_message_history_old_rows o ON o.id = n.id
|
||||
WHERE o IS DISTINCT FROM n
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
DROP INDEX IF EXISTS idx_chat_diff_statuses_pr_title_fts;
|
||||
|
||||
DROP INDEX IF EXISTS idx_chats_title_fts;
|
||||
|
||||
DROP INDEX IF EXISTS idx_chat_messages_search_tsv_pending;
|
||||
|
||||
DROP INDEX IF EXISTS idx_chat_messages_search_tsv;
|
||||
|
||||
ALTER TABLE chat_messages DROP COLUMN IF EXISTS search_tsv;
|
||||
|
||||
DROP FUNCTION IF EXISTS chat_message_search_text(jsonb);
|
||||
@@ -0,0 +1,97 @@
|
||||
CREATE FUNCTION chat_message_search_text(content jsonb) RETURNS text
|
||||
LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
|
||||
SELECT CASE WHEN jsonb_typeof(content) = 'array' THEN (
|
||||
SELECT string_agg(part->>'text', ' ' ORDER BY ordinality)
|
||||
FROM jsonb_array_elements(content) WITH ORDINALITY AS t(part, ordinality)
|
||||
WHERE part->>'type' = 'text'
|
||||
) END
|
||||
$$;
|
||||
|
||||
COMMENT ON FUNCTION chat_message_search_text IS 'Extracts searchable content from chat_messages. Returns NULL for scalar JSON strings (content_version=0). Immutable as it is used in indexes.';
|
||||
|
||||
-- Populated by a background sweep, not at insert time. NULL means pending.
|
||||
ALTER TABLE chat_messages ADD COLUMN search_tsv tsvector;
|
||||
|
||||
COMMENT ON COLUMN chat_messages.search_tsv IS 'Used for full text search. NULL initially, populated async via background job.';
|
||||
|
||||
CREATE INDEX idx_chat_messages_search_tsv ON chat_messages
|
||||
USING GIN (search_tsv)
|
||||
WHERE ((search_tsv IS NOT NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
|
||||
|
||||
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for full text search. Only defined over ''searchable'' rows of chat_messages.';
|
||||
|
||||
CREATE INDEX idx_chat_messages_search_tsv_pending ON chat_messages USING btree (id DESC)
|
||||
WHERE ((search_tsv IS NULL) AND (deleted = false) AND (visibility = ANY (ARRAY['user'::chat_message_visibility, 'both'::chat_message_visibility])) AND (role = ANY (ARRAY['user'::chat_message_role, 'assistant'::chat_message_role])));
|
||||
|
||||
COMMENT ON INDEX idx_chat_messages_search_tsv IS 'Partial index over chat_messages used for populating search_tsv in the background. Only defined over ''searchable'' rows of chat_messages where search_tsv is NULL.';
|
||||
|
||||
CREATE INDEX idx_chats_title_fts ON chats USING GIN (to_tsvector('simple', title));
|
||||
|
||||
COMMENT ON index idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
|
||||
|
||||
CREATE INDEX idx_chat_diff_statuses_pr_title_fts ON chat_diff_statuses USING GIN (to_tsvector('simple', pull_request_title));
|
||||
|
||||
COMMENT ON index idx_chats_title_fts IS 'Used for full text search. Defined over all rows of the chats table.';
|
||||
|
||||
CREATE OR REPLACE FUNCTION set_chat_message_revision_before()
|
||||
RETURNS trigger AS $$
|
||||
DECLARE
|
||||
chat_snapshot_version bigint;
|
||||
cmp chat_messages;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
IF TG_OP = 'UPDATE' THEN
|
||||
IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN
|
||||
RAISE EXCEPTION 'chat_messages.chat_id is immutable';
|
||||
END IF;
|
||||
|
||||
IF OLD.revision IS DISTINCT FROM NEW.revision THEN
|
||||
RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger';
|
||||
END IF;
|
||||
|
||||
cmp := NEW;
|
||||
cmp.search_tsv := OLD.search_tsv;
|
||||
IF OLD IS NOT DISTINCT FROM cmp THEN
|
||||
RETURN NEW;
|
||||
END IF;
|
||||
END IF;
|
||||
|
||||
SELECT snapshot_version INTO chat_snapshot_version
|
||||
FROM chats WHERE id = NEW.chat_id;
|
||||
|
||||
IF chat_snapshot_version IS NULL THEN
|
||||
RAISE EXCEPTION 'chat % does not exist', NEW.chat_id;
|
||||
END IF;
|
||||
|
||||
NEW.revision = chat_snapshot_version;
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
COMMENT ON FUNCTION set_chat_message_revision_before IS 'Component of chatd. Updates chat_snapshot_version when any fields of chat_messages change. Excludes changes to search_tsv as it is not relevant to chatd''s processing loop.';
|
||||
|
||||
CREATE OR REPLACE FUNCTION update_chat_history_after_message_update()
|
||||
RETURNS trigger AS $$
|
||||
BEGIN
|
||||
UPDATE chats c
|
||||
SET history_version = c.snapshot_version,
|
||||
generation_attempt = 0
|
||||
FROM (
|
||||
SELECT DISTINCT n.chat_id
|
||||
FROM chat_message_history_new_rows n
|
||||
JOIN chat_message_history_old_rows o ON o.id = n.id
|
||||
WHERE (to_jsonb(o) - 'search_tsv') IS DISTINCT FROM (to_jsonb(n) - 'search_tsv')
|
||||
) AS affected
|
||||
WHERE c.id = affected.chat_id
|
||||
AND (
|
||||
c.history_version IS DISTINCT FROM c.snapshot_version
|
||||
OR c.generation_attempt <> 0
|
||||
);
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$ LANGUAGE plpgsql;
|
||||
|
||||
COMMENT ON FUNCTION update_chat_history_after_message_update IS 'Component of chatd. Updates history_version and generation_attempt on chats when chat_messages is updated. Excludes changes to search_tsv.';
|
||||
@@ -19,11 +19,13 @@ import (
|
||||
"github.com/golang-migrate/migrate/v4/source/stub"
|
||||
"github.com/google/uuid"
|
||||
"github.com/lib/pq"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/goleak"
|
||||
"golang.org/x/sync/errgroup"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/database/migrations"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
@@ -1866,3 +1868,244 @@ func TestMigration000498SoftDeleteStaleWorkspaceAgents(t *testing.T) {
|
||||
// TestSoftDeleteWorkspaceAgentsByWorkspaceID, plus integration tests
|
||||
// under coderd/coderd_test.go; not retested here.
|
||||
}
|
||||
|
||||
func TestMigration000543ChatMessageSearchText(t *testing.T) {
|
||||
t.Parallel()
|
||||
if testing.Short() {
|
||||
t.SkipNow()
|
||||
}
|
||||
|
||||
sqlDB := testSQLDB(t)
|
||||
require.NoError(t, migrations.Up(sqlDB))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
content sql.NullString
|
||||
want sql.NullString
|
||||
}{
|
||||
{
|
||||
name: "SingleTextPart",
|
||||
content: sql.NullString{String: `[{"type":"text","text":"hello world"}]`, Valid: true},
|
||||
want: sql.NullString{String: "hello world", Valid: true},
|
||||
},
|
||||
{
|
||||
name: "TextInterleavedWithNonText",
|
||||
content: sql.NullString{String: `[
|
||||
{"type":"text","text":"first"},
|
||||
{"type":"reasoning","text":"thinking"},
|
||||
{"type":"tool-call","toolName":"execute"},
|
||||
{"type":"text","text":"second"}
|
||||
]`, Valid: true},
|
||||
want: sql.NullString{String: "first second", Valid: true},
|
||||
},
|
||||
{
|
||||
name: "OnlyNonTextParts",
|
||||
content: sql.NullString{String: `[{"type":"reasoning","text":"thinking"}]`, Valid: true},
|
||||
want: sql.NullString{},
|
||||
},
|
||||
{
|
||||
name: "ScalarContent",
|
||||
content: sql.NullString{String: `"hello"`, Valid: true},
|
||||
want: sql.NullString{},
|
||||
},
|
||||
{
|
||||
name: "EmptyArray",
|
||||
content: sql.NullString{String: `[]`, Valid: true},
|
||||
want: sql.NullString{},
|
||||
},
|
||||
{
|
||||
name: "NullInput",
|
||||
content: sql.NullString{},
|
||||
want: sql.NullString{},
|
||||
},
|
||||
{
|
||||
name: "ElementsMissingTypeOrText",
|
||||
content: sql.NullString{String: `[{"text":"no type"},{"type":"text"},{"type":"text","text":"kept"}]`, Valid: true},
|
||||
want: sql.NullString{String: "kept", Valid: true},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
var got sql.NullString
|
||||
err := sqlDB.QueryRowContext(ctx,
|
||||
`SELECT chat_message_search_text($1::jsonb)`, tc.content,
|
||||
).Scan(&got)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Shared eligibility predicate of the two partial chat_messages search
|
||||
// indexes. Queries must repeat it verbatim.
|
||||
const eligibilityPredicate = `deleted = false
|
||||
AND visibility IN ('user', 'both')
|
||||
AND role IN ('user', 'assistant')`
|
||||
|
||||
func TestMigration000543ChatSearchSchemaIndexes(t *testing.T) {
|
||||
t.Parallel()
|
||||
if testing.Short() {
|
||||
t.SkipNow()
|
||||
}
|
||||
|
||||
sqlDB := testSQLDB(t)
|
||||
require.NoError(t, migrations.Up(sqlDB))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
table string
|
||||
partial bool
|
||||
}{
|
||||
{name: "idx_chat_messages_search_tsv", table: "chat_messages", partial: true},
|
||||
{name: "idx_chat_messages_search_tsv_pending", table: "chat_messages", partial: true},
|
||||
{name: "idx_chats_title_fts", table: "chats", partial: false},
|
||||
{name: "idx_chat_diff_statuses_pr_title_fts", table: "chat_diff_statuses", partial: false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
var table string
|
||||
var partial bool
|
||||
err := sqlDB.QueryRowContext(ctx, `
|
||||
SELECT i.tablename, x.indpred IS NOT NULL
|
||||
FROM pg_indexes i
|
||||
JOIN pg_class c ON c.relname = i.indexname
|
||||
JOIN pg_index x ON x.indexrelid = c.oid
|
||||
WHERE i.indexname = $1`, tc.name,
|
||||
).Scan(&table, &partial)
|
||||
require.NoError(t, err, "index %s should exist", tc.name)
|
||||
require.Equal(t, tc.table, table, "index %s table", tc.name)
|
||||
require.Equal(t, tc.partial, partial, "index %s partial", tc.name)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigration000543ChatSearchSchemaBehavior(t *testing.T) {
|
||||
t.Parallel()
|
||||
if testing.Short() {
|
||||
t.SkipNow()
|
||||
}
|
||||
|
||||
sqlDB := testSQLDB(t)
|
||||
require.NoError(t, migrations.Up(sqlDB))
|
||||
db := database.New(sqlDB)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
owner := dbgen.User(t, db, database.User{})
|
||||
_ = dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "openai", DisplayName: "OpenAI"})
|
||||
modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
||||
IsDefault: true,
|
||||
})
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: owner.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
})
|
||||
|
||||
newMsg := func(role database.ChatMessageRole, visibility database.ChatMessageVisibility, content string) database.ChatMessage {
|
||||
seed := database.ChatMessage{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true},
|
||||
Role: role,
|
||||
Visibility: visibility,
|
||||
}
|
||||
if content != "" {
|
||||
seed.Content = pqtype.NullRawMessage{RawMessage: []byte(content), Valid: true}
|
||||
}
|
||||
return dbgen.ChatMessage(t, db, seed)
|
||||
}
|
||||
textContent := func(text string) string {
|
||||
return `[{"type":"text","text":"` + text + `"}]`
|
||||
}
|
||||
|
||||
pendingIDs := func(ctx context.Context, limit int) []int64 {
|
||||
rows, err := sqlDB.QueryContext(ctx, `
|
||||
SELECT id FROM chat_messages
|
||||
WHERE search_tsv IS NULL AND `+eligibilityPredicate+`
|
||||
ORDER BY id DESC
|
||||
LIMIT $1`, limit)
|
||||
require.NoError(t, err)
|
||||
defer rows.Close()
|
||||
var ids []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
require.NoError(t, rows.Scan(&id))
|
||||
ids = append(ids, id)
|
||||
}
|
||||
require.NoError(t, rows.Err())
|
||||
return ids
|
||||
}
|
||||
|
||||
// Insert regression: RETURNING * must survive the new column, and new
|
||||
// rows must start with search_tsv NULL so they enter the pending queue.
|
||||
eligibleText := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deploy the search feature"))
|
||||
var tsvIsNull bool
|
||||
err := sqlDB.QueryRowContext(ctx,
|
||||
`SELECT search_tsv IS NULL FROM chat_messages WHERE id = $1`, eligibleText.ID,
|
||||
).Scan(&tsvIsNull)
|
||||
require.NoError(t, err)
|
||||
require.True(t, tsvIsNull, "new rows must have search_tsv NULL")
|
||||
|
||||
eligibleNoText := newMsg(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, `[{"type":"reasoning","text":"thinking"}]`)
|
||||
toolMsg := newMsg(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, textContent("tool output about deploy"))
|
||||
modelOnly := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, textContent("model-only deploy note"))
|
||||
deletedMsg := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("deleted deploy message"))
|
||||
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, deletedMsg.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Only eligible rows appear in the queue, newest first. The tool-role,
|
||||
// model-only, and soft-deleted rows are excluded even though their
|
||||
// search_tsv is NULL.
|
||||
require.Equal(t, []int64{eligibleNoText.ID, eligibleText.ID}, pendingIDs(ctx, 10))
|
||||
|
||||
// Sweep-style UPDATE. The '' sentinel (not NULL) marks no-text rows as
|
||||
// swept; NULL means pending, so COALESCE is what drains them from the
|
||||
// queue.
|
||||
_, err = sqlDB.ExecContext(ctx, `
|
||||
UPDATE chat_messages
|
||||
SET search_tsv = COALESCE(to_tsvector('simple', chat_message_search_text(content)), ''::tsvector)
|
||||
WHERE id = ANY($1)`, pq.Array([]int64{eligibleText.ID, eligibleNoText.ID}))
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, pendingIDs(ctx, 10), "swept rows must leave the queue, including no-text rows")
|
||||
|
||||
// Soft-deleting an unswept row removes it from the queue without a sweep.
|
||||
unswept := newMsg(database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, textContent("unswept deploy row"))
|
||||
require.Equal(t, []int64{unswept.ID}, pendingIDs(ctx, 10))
|
||||
_, err = sqlDB.ExecContext(ctx, `UPDATE chat_messages SET deleted = true WHERE id = $1`, unswept.ID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, pendingIDs(ctx, 10))
|
||||
|
||||
// Search contract: populate search_tsv on every row (including
|
||||
// ineligible ones) and assert the search-index predicate filters them.
|
||||
_, err = sqlDB.ExecContext(ctx, `
|
||||
UPDATE chat_messages
|
||||
SET search_tsv = COALESCE(to_tsvector('simple', chat_message_search_text(content)), ''::tsvector)
|
||||
WHERE chat_id = $1`, chat.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
rows, err := sqlDB.QueryContext(ctx, `
|
||||
SELECT id FROM chat_messages
|
||||
WHERE search_tsv @@ websearch_to_tsquery('simple', $1)
|
||||
AND search_tsv IS NOT NULL
|
||||
AND `+eligibilityPredicate+`
|
||||
ORDER BY id`, "deploy")
|
||||
require.NoError(t, err)
|
||||
defer rows.Close()
|
||||
var matched []int64
|
||||
for rows.Next() {
|
||||
var id int64
|
||||
require.NoError(t, rows.Scan(&id))
|
||||
matched = append(matched, id)
|
||||
}
|
||||
require.NoError(t, rows.Err())
|
||||
require.Equal(t, []int64{eligibleText.ID}, matched,
|
||||
"search must exclude deleted, model-only, and tool-role rows (%d %d %d)",
|
||||
toolMsg.ID, modelOnly.ID, deletedMsg.ID)
|
||||
}
|
||||
|
||||
@@ -786,6 +786,7 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams,
|
||||
arg.PrNumber,
|
||||
arg.RepoQuery,
|
||||
arg.PrTitleQuery,
|
||||
arg.Search,
|
||||
arg.OffsetOpt,
|
||||
arg.LimitOpt,
|
||||
)
|
||||
|
||||
Generated
+2
@@ -5123,6 +5123,8 @@ type ChatMessage struct {
|
||||
Revision int64 `db:"revision" json:"revision"`
|
||||
// Stores the selected effort for the turn triggered by this message.
|
||||
ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"`
|
||||
// Used for full text search. NULL initially, populated async via background job.
|
||||
SearchTsv interface{} `db:"search_tsv" json:"search_tsv"`
|
||||
}
|
||||
|
||||
type ChatModelConfig struct {
|
||||
|
||||
Generated
+9
@@ -67,6 +67,12 @@ type sqlcQuerier interface {
|
||||
// created_at ASC flows through to dbpurge's digest truncation; see
|
||||
// buildDigestData in dbpurge.go for the tradeoff rationale.
|
||||
AutoArchiveInactiveChats(ctx context.Context, arg AutoArchiveInactiveChatsParams) ([]AutoArchiveInactiveChatsRow, error)
|
||||
// Backfills chat_messages.search_tsv for pending rows, newest first.
|
||||
// The WHERE clause must match the predicate of
|
||||
// idx_chat_messages_search_tsv_pending exactly so the partial index
|
||||
// serves this query.
|
||||
// NULL means "pending", empty tsvector means "backfilled, no text".
|
||||
BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error)
|
||||
BackoffChatDiffStatus(ctx context.Context, arg BackoffChatDiffStatusParams) error
|
||||
// Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs.
|
||||
BatchDeleteChatHeartbeats(ctx context.Context, arg BatchDeleteChatHeartbeatsParams) (int64, error)
|
||||
@@ -80,6 +86,9 @@ type sqlcQuerier interface {
|
||||
// Calculates the telemetry summary for a given provider, model, and client
|
||||
// combination for telemetry reporting.
|
||||
CalculateAIBridgeInterceptionsTelemetrySummary(ctx context.Context, arg CalculateAIBridgeInterceptionsTelemetrySummaryParams) (CalculateAIBridgeInterceptionsTelemetrySummaryRow, error)
|
||||
// Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
|
||||
// Used to reject input that would silently match nothing.
|
||||
ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error)
|
||||
ClaimPrebuiltWorkspace(ctx context.Context, arg ClaimPrebuiltWorkspaceParams) (ClaimPrebuiltWorkspaceRow, error)
|
||||
CleanTailnetCoordinators(ctx context.Context) error
|
||||
CleanTailnetLostPeers(ctx context.Context) error
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"net"
|
||||
"slices"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -15344,6 +15345,237 @@ func TestGetChatsFilter(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetChatsSearch(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
store, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
org := dbgen.Organization(t, store, database.Organization{})
|
||||
user := dbgen.User(t, store, database.User{})
|
||||
dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID})
|
||||
|
||||
provider := dbgen.AIProviderWithOptionalKey(t, store, database.AIProvider{
|
||||
Type: database.AIProviderTypeOpenai,
|
||||
}, "test-key")
|
||||
|
||||
modelCfg, err := store.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true},
|
||||
Model: "test-model-" + uuid.NewString(),
|
||||
DisplayName: "Test Model",
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
Enabled: true,
|
||||
IsDefault: true,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 80,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
createRoot := func(title string) database.Chat {
|
||||
t.Helper()
|
||||
chat, err := store.InsertChat(ctx, database.InsertChatParams{
|
||||
OrganizationID: org.ID,
|
||||
Status: database.ChatStatusWaiting,
|
||||
ClientType: database.ChatClientTypeUi,
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: title,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return chat
|
||||
}
|
||||
|
||||
createChild := func(root database.Chat, title string) database.Chat {
|
||||
t.Helper()
|
||||
chat, err := store.InsertChat(ctx, database.InsertChatParams{
|
||||
OrganizationID: org.ID,
|
||||
Status: database.ChatStatusWaiting,
|
||||
ClientType: database.ChatClientTypeUi,
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelCfg.ID,
|
||||
Title: title,
|
||||
ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return chat
|
||||
}
|
||||
|
||||
insertMsg := func(chatID uuid.UUID, role database.ChatMessageRole, visibility database.ChatMessageVisibility, text string) database.ChatMessage {
|
||||
t.Helper()
|
||||
msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||
ChatID: chatID,
|
||||
CreatedBy: []uuid.UUID{user.ID},
|
||||
ModelConfigID: []uuid.UUID{modelCfg.ID},
|
||||
Role: []database.ChatMessageRole{role},
|
||||
Content: []string{`[{"type":"text","text":` + strconv.Quote(text) + `}]`},
|
||||
ContentVersion: []int16{1},
|
||||
Visibility: []database.ChatMessageVisibility{visibility},
|
||||
InputTokens: []int64{0},
|
||||
OutputTokens: []int64{0},
|
||||
TotalTokens: []int64{0},
|
||||
ReasoningTokens: []int64{0},
|
||||
CacheCreationTokens: []int64{0},
|
||||
CacheReadTokens: []int64{0},
|
||||
ContextLimit: []int64{0},
|
||||
Compressed: []bool{false},
|
||||
TotalCostMicros: []int64{0},
|
||||
RuntimeMs: []int64{0},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, msgs, 1)
|
||||
return msgs[0]
|
||||
}
|
||||
|
||||
linkPR := func(chatID uuid.UUID, url, state, prTitle string, prNumber int32, gitRemoteOrigin string) {
|
||||
t.Helper()
|
||||
now := time.Now()
|
||||
_, err := store.UpsertChatDiffStatusReference(ctx, database.UpsertChatDiffStatusReferenceParams{
|
||||
ChatID: chatID,
|
||||
Url: sql.NullString{String: url, Valid: true},
|
||||
GitBranch: "main",
|
||||
GitRemoteOrigin: gitRemoteOrigin,
|
||||
StaleAt: now.Add(time.Hour),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = store.UpsertChatDiffStatus(ctx, database.UpsertChatDiffStatusParams{
|
||||
ChatID: chatID,
|
||||
Url: sql.NullString{String: url, Valid: true},
|
||||
PullRequestState: sql.NullString{String: state, Valid: true},
|
||||
PullRequestTitle: prTitle,
|
||||
PrNumber: sql.NullInt32{Int32: prNumber, Valid: prNumber > 0},
|
||||
Additions: 1,
|
||||
Deletions: 1,
|
||||
ChangedFiles: 1,
|
||||
RefreshedAt: now,
|
||||
StaleAt: now.Add(time.Hour),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
titleChat := createRoot("deploy pipeline alpha")
|
||||
|
||||
archivedChat := createRoot("deploy pipeline beta")
|
||||
|
||||
prTitleChat := createRoot("widget work")
|
||||
linkPR(prTitleChat.ID, "https://github.com/acme/widget/pull/42", "open", "Fix authentication bug", 42, "https://github.com/acme/widget.git")
|
||||
|
||||
mergedChat := createRoot("other work")
|
||||
linkPR(mergedChat.ID, "https://github.com/acme/other-repo/pull/7", "merged", "Fix authentication flow", 7, "https://github.com/acme/other-repo.git")
|
||||
|
||||
msgChat := createRoot("plain one")
|
||||
insertMsg(msgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "kubernetes cluster restart")
|
||||
|
||||
assistantMsgChat := createRoot("plain assistant")
|
||||
insertMsg(assistantMsgChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, "grafana dashboard tuning")
|
||||
|
||||
userVisMsgChat := createRoot("plain uservis")
|
||||
insertMsg(userVisMsgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityUser, "vault token rotation")
|
||||
|
||||
assistantUserVisMsgChat := createRoot("plain assistant uservis")
|
||||
insertMsg(assistantUserVisMsgChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, "redis eviction policy")
|
||||
|
||||
deletedMsgChat := createRoot("plain two")
|
||||
deletedMsg := insertMsg(deletedMsgChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "terraform apply failure")
|
||||
|
||||
childParent := createRoot("plain parent")
|
||||
childChat := createChild(childParent, "plain child")
|
||||
insertMsg(childChat.ID, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, "orchestrator saga")
|
||||
|
||||
ineligibleChat := createRoot("plain three")
|
||||
toolMsg := insertMsg(ineligibleChat.ID, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, "forbidden secret token")
|
||||
modelOnlyMsg := insertMsg(ineligibleChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel, "forbidden secret token")
|
||||
|
||||
// Ineligible rows keep search_tsv NULL after backfill.
|
||||
_, err = store.BackfillChatMessagesSearchTsv(ctx, 1000)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Soft-deleted rows stay excluded even though search_tsv remains
|
||||
// populated.
|
||||
err = store.SoftDeleteChatMessageByID(ctx, deletedMsg.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Inserted after backfill: search_tsv IS NULL, must match nothing.
|
||||
pendingChat := createRoot("plain four")
|
||||
insertMsg(pendingChat.ID, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth, "elasticsearch indexing")
|
||||
|
||||
// Prove role/visibility predicates exclude rows even when search_tsv
|
||||
// is set.
|
||||
_, err = sqlDB.ExecContext(ctx,
|
||||
`UPDATE chat_messages SET search_tsv = to_tsvector('simple', 'forbidden secret token') WHERE id = ANY($1)`,
|
||||
pq.Array([]int64{toolMsg.ID, modelOnlyMsg.ID}))
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = store.ArchiveChatByID(ctx, archivedChat.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
allRootIDs := []uuid.UUID{
|
||||
titleChat.ID, archivedChat.ID, prTitleChat.ID, mergedChat.ID,
|
||||
msgChat.ID, assistantMsgChat.ID, userVisMsgChat.ID,
|
||||
assistantUserVisMsgChat.ID, deletedMsgChat.ID, childParent.ID,
|
||||
ineligibleChat.ID, pendingChat.ID,
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
params database.GetChatsParams
|
||||
want []uuid.UUID
|
||||
}{
|
||||
{"Title/Match", database.GetChatsParams{Search: "pipeline alpha"}, []uuid.UUID{titleChat.ID}},
|
||||
{"Title/CaseInsensitiveMultiWord", database.GetChatsParams{Search: "ALPHA DEPLOY"}, []uuid.UUID{titleChat.ID}},
|
||||
{"Title/AndSemantics", database.GetChatsParams{Search: "deploy nonexistent"}, nil},
|
||||
{"PRTitle/Match", database.GetChatsParams{Search: "authentication"}, []uuid.UUID{prTitleChat.ID, mergedChat.ID}},
|
||||
{"Message/Match", database.GetChatsParams{Search: "kubernetes restart"}, []uuid.UUID{msgChat.ID}},
|
||||
{"Message/AssistantRoleMatch", database.GetChatsParams{Search: "grafana tuning"}, []uuid.UUID{assistantMsgChat.ID}},
|
||||
{"Message/UserVisibilityMatch", database.GetChatsParams{Search: "vault rotation"}, []uuid.UUID{userVisMsgChat.ID}},
|
||||
{"Message/AssistantUserVisibilityMatch", database.GetChatsParams{Search: "redis eviction"}, []uuid.UUID{assistantUserVisMsgChat.ID}},
|
||||
{"PRNumber/Match", database.GetChatsParams{Search: "42"}, []uuid.UUID{prTitleChat.ID}},
|
||||
{"PRNumber/NonNumericNoMatch", database.GetChatsParams{Search: "42abc"}, nil},
|
||||
{"PRNumber/OversizedDigitsNoError", database.GetChatsParams{Search: "1111111111111111111111111"}, nil},
|
||||
{"NoMatch", database.GetChatsParams{Search: "zzzqqq"}, nil},
|
||||
{"Message/PendingBackfillNoMatch", database.GetChatsParams{Search: "elasticsearch"}, nil},
|
||||
{"Message/DeletedNoMatch", database.GetChatsParams{Search: "terraform"}, nil},
|
||||
// Parent also excluded: EXISTS is per-chat, not per-tree.
|
||||
{"Message/ChildNotSurfaced", database.GetChatsParams{Search: "orchestrator saga"}, nil},
|
||||
{"Message/IneligibleMessagesNoMatch", database.GetChatsParams{Search: "forbidden secret"}, nil},
|
||||
{"Composed/ArchivedDefaultIncludesAll", database.GetChatsParams{Search: "deploy pipeline"}, []uuid.UUID{titleChat.ID, archivedChat.ID}},
|
||||
{"Composed/ArchivedFalseExcludes", database.GetChatsParams{Search: "deploy pipeline", Archived: sql.NullBool{Bool: false, Valid: true}}, []uuid.UUID{titleChat.ID}},
|
||||
{"Composed/ArchivedTrueOnly", database.GetChatsParams{Search: "deploy pipeline", Archived: sql.NullBool{Bool: true, Valid: true}}, []uuid.UUID{archivedChat.ID}},
|
||||
{"Composed/SearchAndRepo", database.GetChatsParams{Search: "authentication", RepoQuery: "acme/widget"}, []uuid.UUID{prTitleChat.ID}},
|
||||
{"Composed/SearchAndPRStatus", database.GetChatsParams{Search: "authentication", PullRequestStatuses: []string{"merged"}}, []uuid.UUID{mergedChat.ID}},
|
||||
{"EmptySearch/ReturnsAll", database.GetChatsParams{Search: ""}, allRootIDs},
|
||||
{"WhitespaceSearch/ReturnsNothing", database.GetChatsParams{Search: " "}, nil},
|
||||
{"TabOnlySearch/ReturnsNothing", database.GetChatsParams{Search: "\t\t"}, nil},
|
||||
{"EmptySearch/TitleQueryStillWorks", database.GetChatsParams{Search: "", TitleQuery: "pipeline alpha"}, []uuid.UUID{titleChat.ID}},
|
||||
{"EmptySearch/PRTitleQueryStillWorks", database.GetChatsParams{Search: "", PrTitleQuery: "authentication bug"}, []uuid.UUID{prTitleChat.ID}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
params := tt.params
|
||||
params.OwnedOnly = true
|
||||
params.ViewerID = user.ID
|
||||
|
||||
rows, err := store.GetChats(ctx, params)
|
||||
require.NoError(t, err)
|
||||
|
||||
got := make([]uuid.UUID, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
got = append(got, row.Chat.ID)
|
||||
}
|
||||
|
||||
if tt.want == nil {
|
||||
require.Empty(t, got)
|
||||
} else {
|
||||
require.ElementsMatch(t, tt.want, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatHasUnread(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Generated
+103
-10
@@ -6032,6 +6032,36 @@ func (q *sqlQuerier) AutoArchiveInactiveChats(ctx context.Context, arg AutoArchi
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const backfillChatMessagesSearchTsv = `-- name: BackfillChatMessagesSearchTsv :execrows
|
||||
WITH batch AS (
|
||||
SELECT id FROM chat_messages
|
||||
WHERE search_tsv IS NULL
|
||||
AND deleted = false
|
||||
AND visibility IN ('user', 'both')
|
||||
AND role IN ('user', 'assistant')
|
||||
ORDER BY id DESC
|
||||
LIMIT $1::int
|
||||
)
|
||||
UPDATE chat_messages cm
|
||||
SET search_tsv = COALESCE(
|
||||
to_tsvector('simple', chat_message_search_text(cm.content)),
|
||||
''::tsvector)
|
||||
FROM batch WHERE cm.id = batch.id
|
||||
`
|
||||
|
||||
// Backfills chat_messages.search_tsv for pending rows, newest first.
|
||||
// The WHERE clause must match the predicate of
|
||||
// idx_chat_messages_search_tsv_pending exactly so the partial index
|
||||
// serves this query.
|
||||
// NULL means "pending", empty tsvector means "backfilled, no text".
|
||||
func (q *sqlQuerier) BackfillChatMessagesSearchTsv(ctx context.Context, batchSize int32) (int64, error) {
|
||||
result, err := q.db.ExecContext(ctx, backfillChatMessagesSearchTsv, batchSize)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
const backoffChatDiffStatus = `-- name: BackoffChatDiffStatus :exec
|
||||
UPDATE
|
||||
chat_diff_statuses
|
||||
@@ -6096,6 +6126,19 @@ func (q *sqlQuerier) BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUps
|
||||
return err
|
||||
}
|
||||
|
||||
const chatSearchQueryIsEmpty = `-- name: ChatSearchQueryIsEmpty :one
|
||||
SELECT numnode(websearch_to_tsquery('simple', $1::text)) = 0 AS is_empty
|
||||
`
|
||||
|
||||
// Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
|
||||
// Used to reject input that would silently match nothing.
|
||||
func (q *sqlQuerier) ChatSearchQueryIsEmpty(ctx context.Context, search string) (bool, error) {
|
||||
row := q.db.QueryRowContext(ctx, chatSearchQueryIsEmpty, search)
|
||||
var is_empty bool
|
||||
err := row.Scan(&is_empty)
|
||||
return is_empty, err
|
||||
}
|
||||
|
||||
const countChatQueuedMessages = `-- name: CountChatQueuedMessages :one
|
||||
SELECT COUNT(*)::bigint AS count
|
||||
FROM chat_queued_messages
|
||||
@@ -7400,7 +7443,7 @@ func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatP
|
||||
|
||||
const getChatMessageByID = `-- name: GetChatMessageByID :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7436,6 +7479,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -7525,7 +7569,7 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created
|
||||
|
||||
const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7576,6 +7620,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7592,7 +7637,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes
|
||||
|
||||
const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7646,6 +7691,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7662,7 +7708,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar
|
||||
|
||||
const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7729,6 +7775,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7745,7 +7792,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a
|
||||
|
||||
const getChatMessagesByRevisionForStream = `-- name: GetChatMessagesByRevisionForStream :many
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7795,6 +7842,7 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -7827,7 +7875,7 @@ WITH latest_compressed_summary AS (
|
||||
1
|
||||
)
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -7902,6 +7950,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -8588,6 +8637,46 @@ WHERE
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- websearch_to_tsquery accepts quoted phrases, OR, and -negation;
|
||||
-- the 'simple' config folds case and skips stemming.
|
||||
AND CASE
|
||||
WHEN $16::text != '' THEN (
|
||||
-- Served by idx_chats_title_fts.
|
||||
to_tsvector('simple', chats_expanded.title) @@ websearch_to_tsquery('simple', $16)
|
||||
-- Served by idx_chat_diff_statuses_pr_title_fts.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND to_tsvector('simple', cds.pull_request_title) @@ websearch_to_tsquery('simple', $16)
|
||||
)
|
||||
-- The WHERE clause must repeat the predicate of the partial index
|
||||
-- idx_chat_messages_search_tsv so the planner can use it. Additional
|
||||
-- filters should still be fine.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_messages cm
|
||||
WHERE cm.chat_id = chats_expanded.id
|
||||
AND cm.search_tsv IS NOT NULL
|
||||
AND cm.deleted = false
|
||||
AND cm.visibility IN ('user', 'both')
|
||||
AND cm.role IN ('user', 'assistant')
|
||||
AND cm.search_tsv @@ websearch_to_tsquery('simple', $16)
|
||||
)
|
||||
-- Skip an explicit pr_number lookup unless the search is a valid bigint.
|
||||
OR CASE
|
||||
WHEN $16 ~ '^[0-9]{1,18}$' THEN EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND cds.pr_number IS NOT NULL
|
||||
AND cds.pr_number = $16::bigint
|
||||
)
|
||||
ELSE false
|
||||
END
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- Paginate over root chats only. Children are fetched
|
||||
-- separately via GetChildChatsByParentIDs and embedded under
|
||||
-- each parent. Other callers that need the full set should
|
||||
@@ -8604,11 +8693,11 @@ ORDER BY
|
||||
-chats_expanded.pin_order DESC,
|
||||
chats_expanded.updated_at DESC,
|
||||
chats_expanded.id DESC
|
||||
OFFSET $16
|
||||
OFFSET $17
|
||||
LIMIT
|
||||
-- The chat list is unbounded and expected to grow large.
|
||||
-- Default to 50 to prevent accidental excessively large queries.
|
||||
COALESCE(NULLIF($17 :: int, 0), 50)
|
||||
COALESCE(NULLIF($18 :: int, 0), 50)
|
||||
`
|
||||
|
||||
type GetChatsParams struct {
|
||||
@@ -8627,6 +8716,7 @@ type GetChatsParams struct {
|
||||
PrNumber int32 `db:"pr_number" json:"pr_number"`
|
||||
RepoQuery string `db:"repo_query" json:"repo_query"`
|
||||
PrTitleQuery string `db:"pr_title_query" json:"pr_title_query"`
|
||||
Search string `db:"search" json:"search"`
|
||||
OffsetOpt int32 `db:"offset_opt" json:"offset_opt"`
|
||||
LimitOpt int32 `db:"limit_opt" json:"limit_opt"`
|
||||
}
|
||||
@@ -8653,6 +8743,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha
|
||||
arg.PrNumber,
|
||||
arg.RepoQuery,
|
||||
arg.PrTitleQuery,
|
||||
arg.Search,
|
||||
arg.OffsetOpt,
|
||||
arg.LimitOpt,
|
||||
)
|
||||
@@ -9147,7 +9238,7 @@ func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
|
||||
const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one
|
||||
SELECT
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
FROM
|
||||
chat_messages
|
||||
WHERE
|
||||
@@ -9193,6 +9284,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -9702,7 +9794,7 @@ SELECT
|
||||
NULLIF(UNNEST($18::bigint[]), 0),
|
||||
NULLIF(UNNEST($19::bigint[]), 0)
|
||||
RETURNING
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort
|
||||
id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision, reasoning_effort, search_tsv
|
||||
`
|
||||
|
||||
type InsertChatMessagesParams struct {
|
||||
@@ -9781,6 +9873,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa
|
||||
&i.APIKeyID,
|
||||
&i.Revision,
|
||||
&i.ReasoningEffort,
|
||||
&i.SearchTsv,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
@@ -312,6 +312,32 @@ SET
|
||||
WHERE
|
||||
id = @id::bigint;
|
||||
|
||||
-- name: BackfillChatMessagesSearchTsv :execrows
|
||||
-- Backfills chat_messages.search_tsv for pending rows, newest first.
|
||||
-- The WHERE clause must match the predicate of
|
||||
-- idx_chat_messages_search_tsv_pending exactly so the partial index
|
||||
-- serves this query.
|
||||
WITH batch AS (
|
||||
SELECT id FROM chat_messages
|
||||
WHERE search_tsv IS NULL
|
||||
AND deleted = false
|
||||
AND visibility IN ('user', 'both')
|
||||
AND role IN ('user', 'assistant')
|
||||
ORDER BY id DESC
|
||||
LIMIT @batch_size::int
|
||||
)
|
||||
UPDATE chat_messages cm
|
||||
-- NULL means "pending", empty tsvector means "backfilled, no text".
|
||||
SET search_tsv = COALESCE(
|
||||
to_tsvector('simple', chat_message_search_text(cm.content)),
|
||||
''::tsvector)
|
||||
FROM batch WHERE cm.id = batch.id;
|
||||
|
||||
-- name: ChatSearchQueryIsEmpty :one
|
||||
-- Reports whether search text tokenizes to an empty tsquery (e.g. '!!!').
|
||||
-- Used to reject input that would silently match nothing.
|
||||
SELECT numnode(websearch_to_tsquery('simple', @search::text)) = 0 AS is_empty;
|
||||
|
||||
-- name: GetChatByID :one
|
||||
SELECT *
|
||||
FROM chats_expanded
|
||||
@@ -651,6 +677,46 @@ WHERE
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- websearch_to_tsquery accepts quoted phrases, OR, and -negation;
|
||||
-- the 'simple' config folds case and skips stemming.
|
||||
AND CASE
|
||||
WHEN @search::text != '' THEN (
|
||||
-- Served by idx_chats_title_fts.
|
||||
to_tsvector('simple', chats_expanded.title) @@ websearch_to_tsquery('simple', @search)
|
||||
-- Served by idx_chat_diff_statuses_pr_title_fts.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND to_tsvector('simple', cds.pull_request_title) @@ websearch_to_tsquery('simple', @search)
|
||||
)
|
||||
-- The WHERE clause must repeat the predicate of the partial index
|
||||
-- idx_chat_messages_search_tsv so the planner can use it. Additional
|
||||
-- filters should still be fine.
|
||||
OR EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_messages cm
|
||||
WHERE cm.chat_id = chats_expanded.id
|
||||
AND cm.search_tsv IS NOT NULL
|
||||
AND cm.deleted = false
|
||||
AND cm.visibility IN ('user', 'both')
|
||||
AND cm.role IN ('user', 'assistant')
|
||||
AND cm.search_tsv @@ websearch_to_tsquery('simple', @search)
|
||||
)
|
||||
-- Skip an explicit pr_number lookup unless the search is a valid bigint.
|
||||
OR CASE
|
||||
WHEN @search ~ '^[0-9]{1,18}$' THEN EXISTS (
|
||||
SELECT 1
|
||||
FROM chat_diff_statuses cds
|
||||
WHERE cds.chat_id = chats_expanded.id
|
||||
AND cds.pr_number IS NOT NULL
|
||||
AND cds.pr_number = @search::bigint
|
||||
)
|
||||
ELSE false
|
||||
END
|
||||
)
|
||||
ELSE true
|
||||
END
|
||||
-- Paginate over root chats only. Children are fetched
|
||||
-- separately via GetChildChatsByParentIDs and embedded under
|
||||
-- each parent. Other callers that need the full set should
|
||||
|
||||
Reference in New Issue
Block a user