390 lines
12 KiB
Go
390 lines
12 KiB
Go
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
|
|
schema string
|
|
encoded bool
|
|
}
|
|
|
|
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 + `".`, schema: 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), set_config('jit','off',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()
|
|
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
|
|
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
|
|
}
|
|
if err = tx.Commit(); err != nil {
|
|
return err
|
|
}
|
|
r.encoded = encoded
|
|
return nil
|
|
}
|
|
|
|
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 (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, r.text(&p.Title), r.text(&abstract), r.text(&authors), r.text(&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) > 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]
|
|
}
|
|
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()
|
|
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
|
|
}
|
|
order := "p.id DESC"
|
|
if filters.Sort == "title" {
|
|
order = "p.title_key,p.id"
|
|
}
|
|
args = append(args, filters.PageSize, (filters.Page-1)*filters.PageSize)
|
|
// 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 := r.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(r.text(&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 := 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
|
|
}
|
|
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, r.text(&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, r.text(&source.Label), r.text(&source.ExternalID), r.text(&source.URL), r.text(&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, r.text(&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, r.text(&publication.Venue), &publication.Year, r.text(&publication.Track), r.text(&publication.Status), r.text(&publication.EvidenceURL)); err != nil {
|
|
_ = rows.Close()
|
|
return err
|
|
}
|
|
byID[id].Publications = append(byID[id].Publications, publication)
|
|
}
|
|
err = rows.Err()
|
|
_ = rows.Close()
|
|
return err
|
|
}
|