diff --git a/db/task_archives.go b/db/task_archives.go index bef9f3a..c21c89d 100644 --- a/db/task_archives.go +++ b/db/task_archives.go @@ -1,6 +1,7 @@ package db import ( + "bytes" "context" "database/sql" "encoding/json" @@ -423,11 +424,37 @@ func (d *DB) FailTaskArchiveJob(id int64, activeState string, cause error) error } func queryArchiveRows(q interface { - QueryRow(query string, args ...any) *sql.Row -}, inner string, args ...any) (json.RawMessage, error) { - var raw []byte - err := q.QueryRow(`SELECT COALESCE(jsonb_agg(to_jsonb(row_data)), '[]'::jsonb) FROM (`+inner+`) row_data`, args...).Scan(&raw) - return json.RawMessage(raw), err + Query(query string, args ...any) (*sql.Rows, error) +}, inner string, args ...any) (json.RawMessage, int64, error) { + // Do not aggregate the result in PostgreSQL. A jsonb array has a hard limit + // of 256 MiB for its elements, which large LLM request/response histories can + // exceed even though every individual record is valid. Reading row JSON in + // order also avoids building a second copy of the full table in PostgreSQL. + rows, err := q.Query(`SELECT row_to_json(row_data)::text FROM (`+inner+`) row_data`, args...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + + var output bytes.Buffer + output.WriteByte('[') + var count int64 + for rows.Next() { + var raw []byte + if err := rows.Scan(&raw); err != nil { + return nil, 0, err + } + if count > 0 { + output.WriteByte(',') + } + output.Write(raw) + count++ + } + if err := rows.Err(); err != nil { + return nil, 0, err + } + output.WriteByte(']') + return json.RawMessage(output.Bytes()), count, nil } func rawRowCount(raw json.RawMessage) int64 { @@ -492,12 +519,12 @@ func (d *DB) SnapshotTaskArchive(taskID int64) (*TaskArchiveSnapshot, error) { } counts := make(map[string]int64, len(queries)) for _, query := range queries { - raw, err := queryArchiveRows(tx, query.query, query.args...) + raw, count, err := queryArchiveRows(tx, query.query, query.args...) if err != nil { return nil, fmt.Errorf("snapshot %s: %w", query.name, err) } tables[query.name] = raw - counts[query.name] = rawRowCount(raw) + counts[query.name] = count } assetIDs, exclusiveAssetIDs, hosts, exclusiveHosts, err := archiveAssetMetadata(tx, taskID, expID, tables["assets"]) diff --git a/db/task_archives_test.go b/db/task_archives_test.go index dbe282f..c9a6d18 100644 --- a/db/task_archives_test.go +++ b/db/task_archives_test.go @@ -8,6 +8,38 @@ import ( "time" ) +func TestQueryArchiveRowsStreamsJSON(t *testing.T) { + d, err := Open(testDSN(t)) + if err != nil { + t.Skipf("postgres unavailable (%v) — skipping", err) + } + defer d.Close() + + payload := strings.Repeat("large request/response payload ", 64*1024) + raw, count, err := queryArchiveRows(d, `SELECT * FROM (VALUES +($1::bigint,$2::text),($3::bigint,$4::text)) archived_row(id,body) ORDER BY id`, + int64(1), payload, int64(2), "quoted: \"value\"\nline") + if err != nil { + t.Fatal(err) + } + if count != 2 { + t.Fatalf("row count=%d, want 2", count) + } + var rows []struct { + ID int64 `json:"id"` + Body string `json:"body"` + } + if err := json.Unmarshal(raw, &rows); err != nil { + t.Fatal(err) + } + if len(rows) != 2 { + t.Fatalf("decoded row count=%d, want 2", len(rows)) + } + if rows[0].ID != 1 || rows[0].Body != payload || rows[1].ID != 2 || rows[1].Body != "quoted: \"value\"\nline" { + t.Fatal("streamed rows were reordered or truncated") + } +} + func TestTaskArchiveDatabaseRoundTrip(t *testing.T) { d, err := Open(testDSN(t)) if err != nil {