Prepare read-only PostgreSQL paper queries and optional Core wiring

2026-11-07 02:38:28 +08:00
parent 81cf7d9448
commit 028e453d92
9 changed files with 504 additions and 6 deletions
@@ -0,0 +1,345 @@
+package research
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// The fixed Paper Atlas migration a3b677b preserves business columns as TEXT
+// 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
+}
+
+var schemaName = regexp.MustCompile(`^[a-z][a-z0-9_]{0,35}$`)
+
+func NewPostgresRepository(db *sql.DB, schema string) (*PostgresRepository, error) {
+	if db == nil || !schemaName.MatchString(schema) {
+		return nil, errors.New("invalid research database configuration")
+	}
+	return &PostgresRepository{db: db, prefix: `"` + schema + `".`}, nil
+}
+
+func (r *PostgresRepository) query(value string) string {
+	return strings.ReplaceAll(value, "{schema}", r.prefix)
+}
+
+func (r *PostgresRepository) begin(ctx context.Context) (*sql.Tx, error) {
+	tx, err := r.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelRepeatableRead, ReadOnly: true})
+	if err != nil {
+		return nil, err
+	}
+	_, 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)`)
+	if err != nil {
+		_ = tx.Rollback()
+		return nil, err
+	}
+	return tx, nil
+}
+
+// Check resolves required columns and SELECT permissions without scanning the
+// corpus. The migration publishes the complete schema atomically.
+func (r *PostgresRepository) Check(ctx context.Context) error {
+	tx, err := r.begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+	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
+		FROM {schema}papers p LEFT JOIN {schema}source_records r ON r.paper_id=p.id
+		LEFT JOIN {schema}sources so ON so.id=r.source_id LEFT JOIN {schema}identifiers i ON i.paper_id=p.id
+		LEFT JOIN {schema}paper_tags pt ON pt.paper_id=p.id LEFT JOIN {schema}tags t ON t.id=pt.tag_id
+		LEFT JOIN {schema}submissions s ON s.paper_id=p.id LEFT JOIN {schema}editions e ON e.id=s.edition_id
+		LEFT JOIN {schema}venues v ON v.id=e.venue_id LEFT JOIN {schema}decisions d ON d.submission_id=s.id WHERE false`))
+	if err != nil {
+		return err
+	}
+	if err = rows.Close(); err != nil {
+		return err
+	}
+	return tx.Commit()
+}
+
+func whereClause(filters Filters) (string, []any) {
+	clauses, values := []string{}, []any{}
+	if filters.Query != "" {
+		values = append(values, filters.Query)
+		if filters.SearchIn == "author" {
+			// JSON stores author strings or [family,given,suffix] tuples. Plain
+			// full-text terms search names regardless of tuple display order.
+			clauses = append(clauses, `to_tsvector('simple'::regconfig,p.authors_json) @@ plainto_tsquery('simple'::regconfig,$1)`)
+		} else {
+			clauses = append(clauses, `to_tsvector('english'::regconfig,coalesce(p.title,'') || ' ' || coalesce(p.abstract,'')) @@ plainto_tsquery('english'::regconfig,$1)`)
+		}
+	}
+	if filters.Topic != "" {
+		values = append(values, filters.Topic)
+		clauses = append(clauses, fmt.Sprintf(`EXISTS(SELECT 1 FROM {schema}paper_tags pt
+			JOIN {schema}tags t ON t.id=pt.tag_id
+			WHERE pt.paper_id=p.id AND t.namespace='arxiv:category' AND t.value=$%d)`, len(values)))
+	}
+	if len(clauses) == 0 {
+		return "", values
+	}
+	return " WHERE " + strings.Join(clauses, " AND "), values
+}
+
+type scanner interface{ Scan(...any) error }
+
+func 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 {
+		return p, err
+	}
+	p.ID = strconv.FormatInt(id, 10)
+	if abstract.Valid {
+		p.Abstract = &abstract.String
+	}
+	if published.Valid {
+		p.PublishedAt = &published.String
+	}
+	var err error
+	p.Authors, err = decodeAuthors(authors)
+	return p, err
+}
+
+func decodeAuthors(raw string) ([]string, error) {
+	var entries []json.RawMessage
+	if err := json.Unmarshal([]byte(raw), &entries); err != nil {
+		return nil, errors.New("invalid stored authors")
+	}
+	names := []string{}
+	for _, entry := range entries {
+		var name string
+		if json.Unmarshal(entry, &name) == nil {
+			if name = strings.TrimSpace(name); name != "" {
+				names = append(names, name)
+			}
+			continue
+		}
+		var tuple []string
+		if json.Unmarshal(entry, &tuple) != nil || len(tuple) == 0 || len(tuple) > 3 {
+			return nil, errors.New("unsupported stored author")
+		}
+		if len(tuple) >= 2 {
+			tuple[0], tuple[1] = tuple[1], tuple[0]
+		}
+		parts := []string{}
+		for _, part := range tuple {
+			if part = strings.TrimSpace(part); part != "" {
+				parts = append(parts, part)
+			}
+		}
+		if name = strings.Join(parts, " "); name != "" {
+			names = append(names, name)
+		}
+	}
+	return names, nil
+}
+
+func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, error) {
+	result := Page{Items: []Paper{}, Topics: []string{}, Page: filters.Page, PageSize: filters.PageSize}
+	tx, err := r.begin(ctx)
+	if err != nil {
+		return result, err
+	}
+	defer tx.Rollback()
+	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
+	}
+	order := "p.id DESC"
+	if filters.Sort == "title" {
+		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))
+	rows, err := tx.QueryContext(ctx, r.query(query), args...)
+	if err != nil {
+		return result, err
+	}
+	for rows.Next() {
+		p, err := scanPaper(rows)
+		if err != nil {
+			_ = rows.Close()
+			return result, err
+		}
+		result.Items = append(result.Items, p)
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	if err != nil {
+		return result, err
+	}
+	if err = r.enrich(ctx, tx, result.Items); err != nil {
+		return result, err
+	}
+	rows, err = tx.QueryContext(ctx, r.query(`SELECT value FROM {schema}tags WHERE namespace='arxiv:category' ORDER BY value`))
+	if err != nil {
+		return result, err
+	}
+	for rows.Next() {
+		var topic string
+		if err = rows.Scan(&topic); err != nil {
+			_ = rows.Close()
+			return result, err
+		}
+		result.Topics = append(result.Topics, topic)
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	if err != nil {
+		return result, err
+	}
+	return result, tx.Commit()
+}
+
+func (r *PostgresRepository) Get(ctx context.Context, id int64) (Paper, error) {
+	tx, err := r.begin(ctx)
+	if err != nil {
+		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))
+	if err != nil {
+		return p, err
+	}
+	papers := []Paper{p}
+	if err = r.enrich(ctx, tx, papers); err != nil {
+		return p, err
+	}
+	return papers[0], tx.Commit()
+}
+
+// Batch metadata queries avoid per-paper round trips and share the count/page
+// snapshot. User values are bound; only the validated schema is interpolated.
+func (r *PostgresRepository) enrich(ctx context.Context, tx *sql.Tx, papers []Paper) error {
+	if len(papers) == 0 {
+		return nil
+	}
+	ids := make([]int64, len(papers))
+	byID := map[int64]*Paper{}
+	for i := range papers {
+		id, err := strconv.ParseInt(papers[i].ID, 10, 64)
+		if err != nil {
+			return err
+		}
+		ids[i] = id
+		byID[id] = &papers[i]
+	}
+	rows, err := tx.QueryContext(ctx, r.query(`SELECT DISTINCT pt.paper_id,t.value
+		FROM {schema}paper_tags pt JOIN {schema}tags t ON t.id=pt.tag_id
+		WHERE pt.paper_id=ANY($1) AND t.namespace='arxiv:category' ORDER BY pt.paper_id,t.value`), ids)
+	if err != nil {
+		return err
+	}
+	for rows.Next() {
+		var id int64
+		var topic string
+		if err = rows.Scan(&id, &topic); err != nil {
+			_ = rows.Close()
+			return err
+		}
+		byID[id].Topics = append(byID[id].Topics, topic)
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	if err != nil {
+		return err
+	}
+	rows, err = tx.QueryContext(ctx, r.query(`SELECT r.paper_id,s.name,r.external_id,r.url,r.metadata_json
+		FROM {schema}source_records r JOIN {schema}sources s ON s.id=r.source_id
+		WHERE r.paper_id=ANY($1) ORDER BY r.paper_id,r.id`), ids)
+	if err != nil {
+		return err
+	}
+	for rows.Next() {
+		var id int64
+		var source Source
+		var metadata string
+		if err = rows.Scan(&id, &source.Label, &source.ExternalID, &source.URL, &metadata); err != nil {
+			_ = rows.Close()
+			return err
+		}
+		p := byID[id]
+		switch source.Label {
+		case "arxiv":
+			source.Label = "arXiv"
+		case "openreview":
+			source.Label = "OpenReview"
+		case "pmlr":
+			source.Label = "PMLR"
+		}
+		p.Sources = append(p.Sources, source)
+		var fields map[string]json.RawMessage
+		if json.Unmarshal([]byte(metadata), &fields) == nil {
+			for key, target := range map[string]**string{"doi": &p.DOI, "journal-ref": &p.Journal, "comments": &p.Comments} {
+				var value string
+				if *target == nil && json.Unmarshal(fields[key], &value) == nil && strings.TrimSpace(value) != "" {
+					*target = &value
+				}
+			}
+		}
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	if err != nil {
+		return err
+	}
+	rows, err = tx.QueryContext(ctx, r.query(`SELECT paper_id,value FROM {schema}identifiers WHERE paper_id=ANY($1) AND scheme='doi' ORDER BY paper_id,value`), ids)
+	if err != nil {
+		return err
+	}
+	for rows.Next() {
+		var id int64
+		var value string
+		if err = rows.Scan(&id, &value); err != nil {
+			_ = rows.Close()
+			return err
+		}
+		if byID[id].DOI == nil && strings.TrimSpace(value) != "" {
+			byID[id].DOI = &value
+		}
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	if err != nil {
+		return err
+	}
+	rows, err = tx.QueryContext(ctx, r.query(`SELECT s.paper_id,v.acronym,e.year,e.track,coalesce(d.status,'unknown'),coalesce(d.evidence_url,'')
+		FROM {schema}submissions s JOIN {schema}editions e ON e.id=s.edition_id JOIN {schema}venues v ON v.id=e.venue_id
+		LEFT JOIN LATERAL (SELECT status,evidence_url FROM {schema}decisions WHERE submission_id=s.id
+			ORDER BY observed_at DESC,coalesce(decision_date,observed_at) DESC,id DESC LIMIT 1) d ON true
+		WHERE s.paper_id=ANY($1) ORDER BY s.paper_id,e.year DESC,v.acronym,e.track,s.id`), ids)
+	if err != nil {
+		return err
+	}
+	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 {
+			_ = rows.Close()
+			return err
+		}
+		byID[id].Publications = append(byID[id].Publications, publication)
+	}
+	err = rows.Err()
+	_ = rows.Close()
+	return err
+}
@@ -0,0 +1,77 @@
+package research
+
+import (
+	"database/sql"
+	"encoding/json"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestStoredAuthorShapes(t *testing.T) {
+	for _, test := range []struct {
+		raw  string
+		want []string
+	}{
+		{`[]`, []string{}},
+		{`["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"}},
+	} {
+		got, err := decodeAuthors(test.raw)
+		if err != nil || !reflect.DeepEqual(got, test.want) {
+			t.Fatalf("%s: %#v, %v", test.raw, got, err)
+		}
+		encoded, _ := json.Marshal(got)
+		if string(encoded) == "null" {
+			t.Fatal("missing authors must remain an array")
+		}
+	}
+	for _, raw := range []string{`broken`, `{"unexpected":"object"}`, `[42]`, `[["a","b","c","d"]]`} {
+		if _, err := decodeAuthors(raw); err == nil {
+			t.Fatalf("accepted unsupported authors %s", raw)
+		}
+	}
+}
+
+func TestQueryValuesRemainBound(t *testing.T) {
+	query := `graph'); DROP TABLE papers; -- OR author:*`
+	where, args := whereClause(Filters{Query: query, Topic: "cs.LG", SearchIn: "text"})
+	if strings.Contains(where, query) || !reflect.DeepEqual(args, []any{query, "cs.LG"}) {
+		t.Fatalf("unbound user input: %s %#v", where, args)
+	}
+	if !strings.Contains(where, "plainto_tsquery('english'::regconfig,$1)") || !strings.Contains(where, "t.value=$2") {
+		t.Fatalf("unexpected search contract: %s", where)
+	}
+	where, args = whereClause(Filters{Query: "Ada Smith", SearchIn: "author"})
+	if !strings.Contains(where, "p.authors_json") || strings.Contains(where, "p.abstract") || !reflect.DeepEqual(args, []any{"Ada Smith"}) {
+		t.Fatalf("author search includes unrelated text: %s", where)
+	}
+	where, args = whereClause(Filters{Topic: "cs.AI"})
+	if !strings.Contains(where, "t.value=$1") || !reflect.DeepEqual(args, []any{"cs.AI"}) {
+		t.Fatalf("topic-only binding: %s %#v", where, args)
+	}
+	where, args = whereClause(Filters{})
+	if where != "" || len(args) != 0 {
+		t.Fatal("empty filters changed the corpus")
+	}
+}
+
+func TestSchemaCannotEscapeQualification(t *testing.T) {
+	for _, name := range []string{"", `research";DROP SCHEMA public;--`, "a.b", "Public", strings.Repeat("a", 37)} {
+		if _, err := NewPostgresRepository(&sql.DB{}, name); err == nil {
+			t.Fatalf("accepted schema %q", name)
+		}
+	}
+	if _, err := NewPostgresRepository(nil, "research"); err == nil {
+		t.Fatal("accepted nil database")
+	}
+	r, err := NewPostgresRepository(&sql.DB{}, "research_v2")
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := r.query(`SELECT p.id FROM {schema}papers p JOIN {schema}source_records s ON s.paper_id=p.id`)
+	if got != `SELECT p.id FROM "research_v2".papers p JOIN "research_v2".source_records s ON s.paper_id=p.id` {
+		t.Fatalf("unqualified corpus tables: %s", got)
+	}
+}