Read the published research corpus and avoid expensive paging plans [skip ci]

This commit is contained in:
2026-11-07 15:51:17 +08:00
parent 913bf4292b
commit 1363792bb0
10 changed files with 355 additions and 37 deletions
+2 -1
View File
@@ -198,10 +198,11 @@ func (s *service) serve(w http.ResponseWriter, r *http.Request) {
reply(w, 200, result)
return
}
var state interface{ SQLState() string }
switch {
case errors.Is(err, sql.ErrNoRows):
reject(w, 404, "not_found")
case errors.Is(ctx.Err(), context.DeadlineExceeded), errors.Is(err, context.DeadlineExceeded):
case errors.Is(ctx.Err(), context.DeadlineExceeded), errors.Is(err, context.DeadlineExceeded), errors.As(err, &state) && state.SQLState() == "57014":
reject(w, 504, "research_timeout")
default:
reject(w, 503, "research_unavailable")
+7
View File
@@ -100,6 +100,11 @@ func TestQueryDefaultsAndUnicode(t *testing.T) {
}
}
type databaseFailure string
func (failure databaseFailure) Error() string { return "private database error" }
func (failure databaseFailure) SQLState() string { return string(failure) }
func TestStorageErrorsAreRedacted(t *testing.T) {
for _, test := range []struct {
err error
@@ -108,6 +113,8 @@ func TestStorageErrorsAreRedacted(t *testing.T) {
}{
{sql.ErrNoRows, 404, "not_found"},
{context.DeadlineExceeded, 504, "research_timeout"},
{errors.Join(databaseFailure("57014"), errors.New("private statement")), 504, "research_timeout"},
{databaseFailure("42501"), 503, "research_unavailable"},
{errors.New("postgres://private-user:private-password@db/corpus"), 503, "research_unavailable"},
} {
mux := http.NewServeMux()
+61 -17
View File
@@ -15,8 +15,10 @@ import (
// and BIGINT and creates the English title/abstract GIN expression below.
// No schema creation, indexing or other writes belong to this repository.
type PostgresRepository struct {
db *sql.DB
prefix string
db *sql.DB
prefix string
schema string
encoded bool
}
var schemaName = regexp.MustCompile(`^[a-z][a-z0-9_]{0,35}$`)
@@ -25,7 +27,7 @@ func NewPostgresRepository(db *sql.DB, schema string) (*PostgresRepository, erro
if db == nil || !schemaName.MatchString(schema) {
return nil, errors.New("invalid research database configuration")
}
return &PostgresRepository{db: db, prefix: `"` + schema + `".`}, nil
return &PostgresRepository{db: db, prefix: `"` + schema + `".`, schema: schema}, nil
}
func (r *PostgresRepository) query(value string) string {
@@ -39,7 +41,7 @@ func (r *PostgresRepository) begin(ctx context.Context) (*sql.Tx, error) {
}
_, err = tx.ExecContext(ctx, `SELECT set_config('statement_timeout','8000',true),
set_config('lock_timeout','1000',true), set_config('work_mem','2MB',true),
set_config('max_parallel_workers_per_gather','0',true)`)
set_config('max_parallel_workers_per_gather','0',true), set_config('jit','off',true)`)
if err != nil {
_ = tx.Rollback()
return nil, err
@@ -55,6 +57,16 @@ func (r *PostgresRepository) Check(ctx context.Context) error {
return err
}
defer tx.Rollback()
var annotation sql.NullString
err = tx.QueryRowContext(ctx, `SELECT pg_catalog.obj_description(oid,'pg_namespace')
FROM pg_catalog.pg_namespace WHERE nspname=$1`, r.schema).Scan(&annotation)
if err != nil {
return err
}
encoded, err := encodedSchema(annotation)
if err != nil {
return err
}
rows, err := tx.QueryContext(ctx, r.query(`SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published,p.title_key,
r.external_id,r.url,r.metadata_json,so.name,i.scheme,i.value,t.namespace,t.value,
s.id,e.year,e.track,v.acronym,d.status,d.evidence_url,d.observed_at,d.decision_date
@@ -69,7 +81,11 @@ func (r *PostgresRepository) Check(ctx context.Context) error {
if err = rows.Close(); err != nil {
return err
}
return tx.Commit()
if err = tx.Commit(); err != nil {
return err
}
r.encoded = encoded
return nil
}
func whereClause(filters Filters) (string, []any) {
@@ -98,12 +114,12 @@ func whereClause(filters Filters) (string, []any) {
type scanner interface{ Scan(...any) error }
func scanPaper(row scanner) (Paper, error) {
func (r *PostgresRepository) scanPaper(row scanner) (Paper, error) {
p := Paper{Authors: []string{}, Topics: []string{}, Sources: []Source{}}
var id int64
var abstract, published sql.NullString
var authors string
if err := row.Scan(&id, &p.Title, &abstract, &authors, &published); err != nil {
if err := row.Scan(&id, r.text(&p.Title), r.text(&abstract), r.text(&authors), r.text(&published)); err != nil {
return p, err
}
p.ID = strconv.FormatInt(id, 10)
@@ -133,9 +149,14 @@ func decodeAuthors(raw string) ([]string, error) {
continue
}
var tuple []string
if json.Unmarshal(entry, &tuple) != nil || len(tuple) == 0 || len(tuple) > 3 {
if json.Unmarshal(entry, &tuple) != nil || len(tuple) == 0 || len(tuple) > 4 {
return nil, errors.New("unsupported stored author")
}
// arXiv may append an affiliation after family, given and suffix.
// It is source metadata, not part of the display name.
if len(tuple) == 4 {
tuple = tuple[:3]
}
if len(tuple) >= 2 {
tuple[0], tuple[1] = tuple[1], tuple[0]
}
@@ -159,6 +180,10 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
return result, err
}
defer tx.Rollback()
if r.encoded {
filters.Query = encodeText(filters.Query)
filters.Topic = encodeText(filters.Topic)
}
where, args := whereClause(filters)
if err = tx.QueryRowContext(ctx, r.query(`SELECT count(*) FROM {schema}papers p`+where), args...).Scan(&result.Total); err != nil {
return result, err
@@ -168,14 +193,33 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
order = "p.title_key,p.id"
}
args = append(args, filters.PageSize, (filters.Page-1)*filters.PageSize)
query := fmt.Sprintf(`SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published
FROM {schema}papers p%s ORDER BY %s LIMIT $%d OFFSET $%d`, where, order, len(args)-1, len(args))
// Select IDs before loading potentially TOASTed titles/abstracts. With
// sparse categories, selecting full rows here can make the planner walk
// hundreds of thousands of wide heap rows before finding eight matches.
candidateQuery := `SELECT p.id FROM {schema}papers p`
candidateOrder := order
prefix := ""
if where != "" {
columns := "p.id"
if filters.Sort == "title" {
columns += ",p.title_key"
}
// Apply search before LIMIT/ORDER BY so a sparse term cannot turn
// into a backward heap scan that recomputes vectors for millions.
prefix = fmt.Sprintf(`matches AS MATERIALIZED (SELECT %s FROM {schema}papers p%s), `, columns, where)
candidateQuery = "SELECT id FROM matches"
candidateOrder = strings.ReplaceAll(order, "p.", "")
}
query := fmt.Sprintf(`WITH %spage_ids AS MATERIALIZED (
%s ORDER BY %s LIMIT $%d OFFSET $%d
) SELECT p.id,p.title,p.abstract,p.authors_json,p.first_published
FROM page_ids JOIN {schema}papers p ON p.id=page_ids.id ORDER BY %s`, prefix, candidateQuery, candidateOrder, len(args)-1, len(args), order)
rows, err := tx.QueryContext(ctx, r.query(query), args...)
if err != nil {
return result, err
}
for rows.Next() {
p, err := scanPaper(rows)
p, err := r.scanPaper(rows)
if err != nil {
_ = rows.Close()
return result, err
@@ -196,7 +240,7 @@ func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, e
}
for rows.Next() {
var topic string
if err = rows.Scan(&topic); err != nil {
if err = rows.Scan(r.text(&topic)); err != nil {
_ = rows.Close()
return result, err
}
@@ -216,7 +260,7 @@ func (r *PostgresRepository) Get(ctx context.Context, id int64) (Paper, error) {
return Paper{}, err
}
defer tx.Rollback()
p, err := scanPaper(tx.QueryRowContext(ctx, r.query(`SELECT id,title,abstract,authors_json,first_published FROM {schema}papers WHERE id=$1`), id))
p, err := r.scanPaper(tx.QueryRowContext(ctx, r.query(`SELECT id,title,abstract,authors_json,first_published FROM {schema}papers WHERE id=$1`), id))
if err != nil {
return p, err
}
@@ -252,7 +296,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var topic string
if err = rows.Scan(&id, &topic); err != nil {
if err = rows.Scan(&id, r.text(&topic)); err != nil {
_ = rows.Close()
return err
}
@@ -273,7 +317,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
var id int64
var source Source
var metadata string
if err = rows.Scan(&id, &source.Label, &source.ExternalID, &source.URL, &metadata); err != nil {
if err = rows.Scan(&id, r.text(&source.Label), r.text(&source.ExternalID), r.text(&source.URL), r.text(&metadata)); err != nil {
_ = rows.Close()
return err
}
@@ -309,7 +353,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var value string
if err = rows.Scan(&id, &value); err != nil {
if err = rows.Scan(&id, r.text(&value)); err != nil {
_ = rows.Close()
return err
}
@@ -333,7 +377,7 @@ func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Pa
for rows.Next() {
var id int64
var publication Publication
if err = rows.Scan(&id, &publication.Venue, &publication.Year, &publication.Track, &publication.Status, &publication.EvidenceURL); err != nil {
if err = rows.Scan(&id, r.text(&publication.Venue), &publication.Year, r.text(&publication.Track), r.text(&publication.Status), r.text(&publication.EvidenceURL)); err != nil {
_ = rows.Close()
return err
}
+36 -4
View File
@@ -102,7 +102,7 @@ func TestPostgresMetadataPreservesEvidence(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if len(p.Sources) != 2 || p.DOI == nil || *p.DOI != "10.fixture/other" || len(p.Publications) != 1 || p.Publications[0].Status != "unknown" || p.Publications[0].EvidenceURL != "" {
if !reflect.DeepEqual(p.Authors, []string{"Carla Network"}) || len(p.Sources) != 2 || p.DOI == nil || *p.DOI != "10.fixture/other" || len(p.Publications) != 1 || p.Publications[0].Status != "unknown" || p.Publications[0].EvidenceURL != "" {
t.Fatalf("unknown decision/source mapping: %+v", p)
}
p, err = r.Get(ctx, 2)
@@ -120,16 +120,19 @@ func TestPostgresCorpusIsReadOnlyAndHTTPWorks(t *testing.T) {
if err != nil {
t.Fatal(err)
}
var readOnly, isolation string
var readOnly, isolation, jit string
if err = tx.QueryRowContext(ctx, "SHOW transaction_read_only").Scan(&readOnly); err != nil {
t.Fatal(err)
}
if err = tx.QueryRowContext(ctx, "SHOW transaction_isolation").Scan(&isolation); err != nil {
t.Fatal(err)
}
if err = tx.QueryRowContext(ctx, "SHOW jit").Scan(&jit); err != nil {
t.Fatal(err)
}
_ = tx.Rollback()
if readOnly != "on" || isolation != "repeatable read" {
t.Fatalf("unsafe read transaction: %s %s", readOnly, isolation)
if readOnly != "on" || isolation != "repeatable read" || jit != "off" {
t.Fatalf("unexpected read transaction settings: %s %s jit=%s", readOnly, isolation, jit)
}
_, err = r.db.ExecContext(ctx, `UPDATE research_fixture.papers SET title='unexpected write' WHERE id=1`)
var state interface{ SQLState() string }
@@ -152,3 +155,32 @@ func TestPostgresCorpusIsReadOnlyAndHTTPWorks(t *testing.T) {
t.Fatalf("missing HTTP status %d", w.Code)
}
}
func TestPostgresPublishedTextCodec(t *testing.T) {
legacy, ctx := integrationRepository(t)
r, err := NewPostgresRepository(legacy.db, "research_encoded_fixture")
if err != nil {
t.Fatal(err)
}
if err = r.Check(ctx); err != nil {
t.Fatal(err)
}
p, err := r.Get(ctx, 1)
if err != nil {
t.Fatal(err)
}
if p.Title != "Graph\x00original\x01" || p.Abstract == nil || *p.Abstract != "Text\x00tail" || !reflect.DeepEqual(p.Authors, []string{"Ada\x00Smith", "Group\x01Name"}) || p.Comments == nil || *p.Comments != "Original\x00comment" {
t.Fatalf("published text was not restored: %+v", p)
}
mux := http.NewServeMux()
Register(mux, r, map[string]string{"thesisly": "http://localhost:17322/auth/callback"})
w := request(mux, "GET", "/projects/thesisly/research/papers/1")
var decoded Paper
if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &decoded) != nil || decoded.Title != p.Title || decoded.Abstract == nil || *decoded.Abstract != *p.Abstract {
t.Fatalf("HTTP lost original control characters: %d %s", w.Code, w.Body.String())
}
page, err := r.List(ctx, Filters{Page: 1, PageSize: 8, Query: "graph"})
if err != nil || page.Total != 2 || len(page.Items) != 2 || page.Items[1].Title != p.Title {
t.Fatalf("encoded corpus search: %+v %v", page, err)
}
}
+2 -1
View File
@@ -14,6 +14,7 @@ func TestStoredAuthorShapes(t *testing.T) {
want []string
}{
{`[]`, []string{}},
{`[["Buhrman","Harry","","CWI"],["Family","Given","Jr.","University"]]`, []string{"Harry Buhrman", "Given Family Jr."}},
{`["The BABAR Collaboration","B. Aubert"]`, []string{"The BABAR Collaboration", "B. Aubert"}},
{`[["Balázs","C.",""],["Berger","E. L.",""],["Smith","Ada","Jr."]]`, []string{"C. Balázs", "E. L. Berger", "Ada Smith Jr."}},
{`[" 张三 ",["Only family"],["","Only given",""],""]`, []string{"张三", "Only family", "Only given"}},
@@ -27,7 +28,7 @@ func TestStoredAuthorShapes(t *testing.T) {
t.Fatal("missing authors must remain an array")
}
}
for _, raw := range []string{`broken`, `{"unexpected":"object"}`, `[42]`, `[["a","b","c","d"]]`} {
for _, raw := range []string{`broken`, `{"unexpected":"object"}`, `[42]`, `[["a","b","c","d","e"]]`} {
if _, err := decodeAuthors(raw); err == nil {
t.Fatalf("accepted unsupported authors %s", raw)
}
+23 -1
View File
@@ -33,7 +33,7 @@ CREATE INDEX decisions_submission_order ON research_fixture.decisions(submission
INSERT INTO research_fixture.papers VALUES
(1,'Graph & learning <original>','Neural graph models preserve original abstracts.','graph learning','[["Smith","Ada",""],"The Research Group"]','2024-01-02'),
(2,'Alpha topology',NULL,'alpha topology','["Bruno Park"]',NULL),
(3,'Graph systems','Graph systems and data.','graph systems','["Carla Network"]','2023-01-01'),
(3,'Graph systems','Graph systems and data.','graph systems','[["Network","Carla","","Fixture University"]]','2023-01-01'),
(4,'Zebra methods','Evidence and uncertainty.','zebra methods','[]','2022-01-01'),
(5,'Alpha Topology','Another independent record.','alpha topology','["Dana Other"]','2025-01-01');
INSERT INTO research_fixture.sources VALUES(10,'arxiv'),(20,'pmlr');
@@ -57,3 +57,25 @@ CREATE ROLE foundry_research_fixture LOGIN;
REVOKE ALL ON SCHEMA research_fixture FROM PUBLIC;
GRANT USAGE ON SCHEMA research_fixture TO foundry_research_fixture;
GRANT SELECT ON ALL TABLES IN SCHEMA research_fixture TO foundry_research_fixture;
-- Separate encoded schema mirrors the published Paper Atlas nul-soh-v1 format.
CREATE SCHEMA research_encoded_fixture;
COMMENT ON SCHEMA research_encoded_fixture IS 'paper-atlas:text-codec=nul-soh-v1';
DO $$
DECLARE item text;
BEGIN
FOREACH item IN ARRAY ARRAY['papers','sources','source_records','identifiers','tags','paper_tags','venues','editions','submissions','decisions'] LOOP
EXECUTE format('CREATE TABLE research_encoded_fixture.%I (LIKE research_fixture.%I INCLUDING ALL)',item,item);
EXECUTE format('INSERT INTO research_encoded_fixture.%I SELECT * FROM research_fixture.%I',item,item);
END LOOP;
END $$;
UPDATE research_encoded_fixture.papers
SET title='Graph'||chr(1)||chr(1)||'original'||chr(1)||chr(2),
abstract='Text'||chr(1)||chr(1)||'tail',
authors_json='["Ada\u0000Smith","Group\u0001Name"]'
WHERE id=1;
UPDATE research_encoded_fixture.source_records
SET metadata_json='{"comments":"Original\u0000comment"}' WHERE id=101;
REVOKE ALL ON SCHEMA research_encoded_fixture FROM PUBLIC;
GRANT USAGE ON SCHEMA research_encoded_fixture TO foundry_research_fixture;
GRANT SELECT ON ALL TABLES IN SCHEMA research_encoded_fixture TO foundry_research_fixture;
+88
View File
@@ -0,0 +1,88 @@
package research
import (
"database/sql"
"errors"
"strings"
)
// Matches Paper Atlas repair 4e62710's nul-soh-v1 protocol: NUL -> SOH/SOH and literal
// SOH -> SOH/STX. All other bytes retain their original binary ordering.
const textCodec = "nul-soh-v1"
func encodedSchema(annotation sql.NullString) (bool, error) {
const prefix = "paper-atlas:text-codec="
if !annotation.Valid || !strings.HasPrefix(annotation.String, prefix) {
return false, nil
}
if strings.TrimPrefix(annotation.String, prefix) != textCodec {
return false, errors.New("unsupported research text codec")
}
return true, nil
}
func encodeText(value string) string {
return strings.ReplaceAll(strings.ReplaceAll(value, "\x01", "\x01\x02"), "\x00", "\x01\x01")
}
func decodeText(value string) (string, error) {
if !strings.ContainsRune(value, '\x01') {
return value, nil
}
var result strings.Builder
for {
position := strings.IndexByte(value, 1)
if position < 0 {
result.WriteString(value)
return result.String(), nil
}
result.WriteString(value[:position])
if position+1 >= len(value) {
return "", errors.New("invalid research text escape")
}
switch value[position+1] {
case 1:
result.WriteByte(0)
case 2:
result.WriteByte(1)
default:
return "", errors.New("invalid research text escape")
}
value = value[position+2:]
}
}
type textColumn struct {
encoded bool
target any
}
func (column textColumn) Scan(source any) error {
var value sql.NullString
if err := value.Scan(source); err != nil {
return err
}
if value.Valid && column.encoded {
var err error
value.String, err = decodeText(value.String)
if err != nil {
return err
}
}
switch target := column.target.(type) {
case *string:
if !value.Valid {
return errors.New("unexpected null research text")
}
*target = value.String
case *sql.NullString:
*target = value
default:
return errors.New("unsupported research text destination")
}
return nil
}
func (r *PostgresRepository) text(target any) sql.Scanner {
return textColumn{encoded: r.encoded, target: target}
}
+69
View File
@@ -0,0 +1,69 @@
package research
import (
"database/sql"
"strings"
"testing"
)
func TestCorpusTextCodecRoundTripAndOrder(t *testing.T) {
values := []string{"", "\x00", "\x00\x01", "\x01", "\x01\x02", "\x02", "A", "A\x00", "A\x01", "Ada 张三 Ω", "Z\x00\x01\x02尾"}
for _, value := range values {
encoded := encodeText(value)
decoded, err := decodeText(encoded)
if err != nil || decoded != value || strings.ContainsRune(encoded, 0) {
t.Fatalf("round trip failed for %q", value)
}
for _, other := range values {
if strings.Compare(value, other) != strings.Compare(encoded, encodeText(other)) {
t.Fatalf("binary ordering changed for %q %q", value, other)
}
}
}
for _, invalid := range []string{"\x01", "prefix\x01x", "\x01\x00"} {
if _, err := decodeText(invalid); err == nil {
t.Fatalf("accepted malformed escape %q", invalid)
}
}
}
func TestCorpusTextColumnsRespectCodecAndNulls(t *testing.T) {
var required string
if err := (textColumn{target: &required}).Scan("\x01\x01"); err != nil || required != "\x01\x01" {
t.Fatal("legacy text was decoded")
}
if err := (textColumn{encoded: true, target: &required}).Scan([]byte("before\x01\x01after\x01\x02")); err != nil || required != "before\x00after\x01" {
t.Fatalf("encoded text did not decode: %q %v", required, err)
}
var nullable sql.NullString
if err := (textColumn{encoded: true, target: &nullable}).Scan(nil); err != nil || nullable.Valid {
t.Fatal("null changed into empty text")
}
if err := (textColumn{encoded: true, target: &nullable}).Scan("\x01\x01"); err != nil || !nullable.Valid || nullable.String != "\x00" {
t.Fatal("nullable text was not restored")
}
if err := (textColumn{encoded: true, target: &required}).Scan(nil); err == nil {
t.Fatal("accepted null for required text")
}
if err := (textColumn{encoded: true, target: &required}).Scan("\x01x"); err == nil {
t.Fatal("accepted corrupt encoded text")
}
}
func TestCorpusSchemaCodecMarker(t *testing.T) {
for _, annotation := range []sql.NullString{{}, {String: "unrelated schema description", Valid: true}} {
encoded, err := encodedSchema(annotation)
if encoded || err != nil {
t.Fatal("legacy schema changed encoding")
}
}
encoded, err := encodedSchema(sql.NullString{String: "paper-atlas:text-codec=nul-soh-v1", Valid: true})
if !encoded || err != nil {
t.Fatal("published codec was not recognized")
}
for _, name := range []string{"unknown", "", "legacy"} {
if _, err := encodedSchema(sql.NullString{String: "paper-atlas:text-codec=" + name, Valid: true}); err == nil {
t.Fatal("unknown codec was accepted")
}
}
}