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

Skip to content
Merged
24 changes: 0 additions & 24 deletions coderd/database/dbpurge/dbpurge_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2775,27 +2775,6 @@ func TestDeleteOldChatFiles(t *testing.T) {
now.Add(-10*24*time.Hour), recentArchivedChat.ID)
require.NoError(t, err)

// File F: 31 days old, in BOTH an active chat AND an old archived chat -> should be retained.
fileF := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour))
anotherOldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour))
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: anotherOldArchivedChat.ID,
MaxFileLinks: 100,
FileIds: []uuid.UUID{fileF},
})
require.NoError(t, err)
_, err = rawDB.ExecContext(ctx, "UPDATE chats SET updated_at = $1 WHERE id = $2",
now.Add(-31*24*time.Hour), anotherOldArchivedChat.ID)
require.NoError(t, err)

activeChatForF := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now)
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: activeChatForF.ID,
MaxFileLinks: 100,
FileIds: []uuid.UUID{fileF},
})
require.NoError(t, err)

done := awaitDoTick(ctx, t, clk)
closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk))
defer closer.Close()
Expand All @@ -2806,9 +2785,6 @@ func TestDeleteOldChatFiles(t *testing.T) {

_, err = db.GetChatFileByID(ctx, fileE)
require.NoError(t, err, "file E in recently archived chat should be retained")

_, err = db.GetChatFileByID(ctx, fileF)
require.NoError(t, err, "file F in active + old archived chat should be retained")
},
},
{
Expand Down
5 changes: 3 additions & 2 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
CREATE INDEX idx_chat_file_links_file_id ON chat_file_links (file_id);

ALTER TABLE chat_file_links
DROP CONSTRAINT chat_file_links_file_id_key;
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
-- A file belongs to one chat. Drop any duplicate links, keeping the link on
-- the oldest chat, before enforcing it.
DELETE FROM chat_file_links l
Comment thread
ethanndickson marked this conversation as resolved.
USING chat_file_links o, chats lc, chats oc
WHERE o.file_id = l.file_id
AND lc.id = l.chat_id
AND oc.id = o.chat_id
AND (oc.created_at, oc.id) < (lc.created_at, lc.id);

ALTER TABLE chat_file_links
ADD CONSTRAINT chat_file_links_file_id_key UNIQUE (file_id);

-- The unique index replaces the purge lookup index.
DROP INDEX idx_chat_file_links_file_id;
7 changes: 5 additions & 2 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

109 changes: 109 additions & 0 deletions coderd/database/querier_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2116,6 +2116,115 @@ func TestLinkChatFilesDeduplicatesInput(t *testing.T) {
require.Equal(t, file.ID, files[0].ID)
}

func TestLinkChatFilesEvictsOldest(t *testing.T) {
t.Parallel()
if testing.Short() {
t.SkipNow()
}

ctx := testutil.Context(t, testutil.WaitMedium)
sqlDB := testSQLDB(t)
err := migrations.Up(sqlDB)
require.NoError(t, err)
db := database.New(sqlDB)

user := dbgen.User(t, db, database.User{})
org := dbgen.Organization(t, db, database.Organization{})
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{})
newChat := func() database.Chat {
return dbgen.Chat(t, db, database.Chat{
OrganizationID: org.ID,
OwnerID: user.ID,
LastModelConfigID: model.ID,
})
}
const maxLinks = 3
base := dbtime.Now().Add(-time.Hour)
newFiles := func(n int) []uuid.UUID {
ids := make([]uuid.UUID, 0, n)
for i := 0; i < n; i++ {
file, err := db.InsertChatFile(ctx, database.InsertChatFileParams{
OwnerID: user.ID,
OrganizationID: org.ID,
Name: fmt.Sprintf("file-%d.txt", i),
Mimetype: "text/plain",
Data: []byte("data"),
})
require.NoError(t, err)
_, err = sqlDB.ExecContext(ctx,
"UPDATE chat_files SET created_at = $1 WHERE id = $2",
base.Add(time.Duration(i)*time.Second), file.ID)
require.NoError(t, err)
ids = append(ids, file.ID)
}
return ids
}
linkedIDs := func(chatID uuid.UUID) []uuid.UUID {
files, err := db.GetChatFileMetadataByChatID(ctx, chatID)
require.NoError(t, err)
ids := make([]uuid.UUID, 0, len(files))
for _, f := range files {
ids = append(ids, f.ID)
}
return ids
}

// Linking one file past the cap deletes the oldest file.
chat := newChat()
files := newFiles(maxLinks + 1)
rejected, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: chat.ID,
FileIds: files[:maxLinks],
MaxFileLinks: maxLinks,
})
require.NoError(t, err)
require.Zero(t, rejected)
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: chat.ID,
FileIds: files[maxLinks:],
MaxFileLinks: maxLinks,
})
require.NoError(t, err)
require.Zero(t, rejected)
require.Equal(t, files[1:], linkedIDs(chat.ID))
_, err = db.GetChatFileByID(ctx, files[0])
require.ErrorIs(t, err, sql.ErrNoRows)

// A single batch over the cap is rejected and deletes nothing.
chat = newChat()
files = newFiles(maxLinks + 1)
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: chat.ID,
FileIds: files,
MaxFileLinks: maxLinks,
})
require.NoError(t, err)
require.EqualValues(t, maxLinks+1, rejected)
require.Empty(t, linkedIDs(chat.ID))
for _, id := range files {
_, err = db.GetChatFileByID(ctx, id)
require.NoError(t, err)
}

// A file links to one chat only.
chat = newChat()
files = newFiles(1)
rejected, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: chat.ID,
FileIds: files,
MaxFileLinks: maxLinks,
})
require.NoError(t, err)
require.Zero(t, rejected)
_, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{
ChatID: newChat().ID,
FileIds: files,
MaxFileLinks: maxLinks,
})
require.True(t, database.IsUniqueViolation(err, database.UniqueChatFileLinksFileIDKey))
require.Equal(t, files, linkedIDs(chat.ID))
}

func TestGetChatFileDataPrefixesByIDs(t *testing.T) {
t.Parallel()
if testing.Short() {
Expand Down
57 changes: 38 additions & 19 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

53 changes: 36 additions & 17 deletions coderd/database/queries/chats.sql
Original file line number Diff line number Diff line change
Expand Up @@ -1755,35 +1755,54 @@ WHERE chat_id = @chat_id::uuid
ORDER BY source ASC;

-- name: LinkChatFilesAfterLock :one
-- LinkChatFilesAfterLock requires the chat row lock.
-- The lock serializes cap checks. The result counts rejected new links.
WITH current AS (
SELECT COUNT(*) AS cnt
FROM chat_file_links
WHERE chat_id = @chat_id::uuid
-- LinkChatFilesAfterLock requires the chat row lock. When the batch would
-- exceed the cap, the oldest files on the chat are deleted to make room; the
-- cascade removes their links. A file links to at most one chat, so no other
-- chat can lose a file here. The batch is rejected only when the batch itself
-- exceeds the cap.
WITH new_links AS (
SELECT DISTINCT unnest(@file_ids::uuid[]) AS file_id
),
new_links AS (
SELECT DISTINCT @chat_id::uuid AS chat_id, unnest(@file_ids::uuid[]) AS file_id
fits AS (
SELECT (SELECT COUNT(*) FROM new_links) <= @max_file_links::int AS ok
),
genuinely_new AS (
SELECT nl.chat_id, nl.file_id
FROM new_links nl
SELECT nl.file_id FROM new_links nl
WHERE NOT EXISTS (
SELECT 1 FROM chat_file_links cfl
WHERE cfl.chat_id = nl.chat_id AND cfl.file_id = nl.file_id
WHERE cfl.chat_id = @chat_id::uuid AND cfl.file_id = nl.file_id
)
),
needed AS (
SELECT GREATEST(
(SELECT COUNT(*) FROM chat_file_links WHERE chat_id = @chat_id::uuid)
+ (SELECT COUNT(*) FROM genuinely_new)
- @max_file_links::int, 0)::int AS n
Comment thread
ethanndickson marked this conversation as resolved.
),
candidates AS (
SELECT cf.id
FROM chat_file_links cfl
JOIN chat_files cf ON cf.id = cfl.file_id
WHERE cfl.chat_id = @chat_id::uuid
AND NOT EXISTS (SELECT 1 FROM new_links nl WHERE nl.file_id = cf.id)
ORDER BY cf.created_at ASC, cf.id ASC
LIMIT (SELECT n FROM needed)
Comment thread
ethanndickson marked this conversation as resolved.
),
evicted AS (
DELETE FROM chat_files cf
USING candidates c
WHERE cf.id = c.id AND (SELECT ok FROM fits)
Comment thread
ethanndickson marked this conversation as resolved.
RETURNING cf.id
),
inserted AS (
INSERT INTO chat_file_links (chat_id, file_id)
SELECT gn.chat_id, gn.file_id
FROM genuinely_new gn, current c
WHERE c.cnt + (SELECT COUNT(*) FROM genuinely_new) <= @max_file_links::int
SELECT @chat_id::uuid, gn.file_id FROM genuinely_new gn
WHERE (SELECT ok FROM fits)
ON CONFLICT (chat_id, file_id) DO NOTHING
RETURNING file_id
)
SELECT
(SELECT COUNT(*)::int FROM genuinely_new) -
(SELECT COUNT(*)::int FROM inserted) AS rejected_new_files;
SELECT (SELECT COUNT(*)::int FROM genuinely_new)
- (SELECT COUNT(*)::int FROM inserted) AS rejected_new_files;

-- name: UpdateChatStatus :one
WITH updated_chat AS (
Expand Down
1 change: 1 addition & 0 deletions coderd/database/unique_constraint.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion coderd/exp_chats.go
Original file line number Diff line number Diff line change
Expand Up @@ -6626,7 +6626,7 @@ func writeChatFileError(ctx context.Context, rw http.ResponseWriter, err error)
case errors.Is(err, chatstate.ErrChatFileCapExceeded):
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Chat attachment limit reached.",
Detail: fmt.Sprintf("A chat can reference at most %d attachments. Remove some attachments or start a new chat.", codersdk.MaxChatFileIDs),
Detail: fmt.Sprintf("A message can include at most %d attachments. Remove some attachments and retry.", codersdk.MaxChatFileIDs),
})
case errors.Is(err, chatstate.ErrChatFileUnavailable):
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Expand Down
Loading
Loading