Files
foundry-core/internal/research/postgres.go
T
codex-bot 028e453d92
Foundry Go checks / check (push) Successful in 3m57s
Foundry Go checks / check (pull_request) Successful in 5m2s
Prepare read-only PostgreSQL paper queries and optional Core wiring
2026-11-07 02:38:28 +08:00

346 lines
11 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
}
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
}