foundry-core/internal/research/postgres.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) positions(ctx context.Context, tx *sql.Tx, query string, args []any) ([]position, error) {
	rows, err := tx.QueryContext(ctx, r.query(query), args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	positions := []position{}
	for rows.Next() {
		var p position
		// Preserve the stored title key, including the published text codec.
		if err = rows.Scan(&p.ID, &p.Key); err != nil {
			return nil, err
		}
		positions = append(positions, p)
	}
	return positions, rows.Err()
}

func (r *PostgresRepository) List(ctx context.Context, filters Filters) (Page, error) {
	result := Page{Items: []Paper{}, Topics: []string{}, PageSize: filters.PageSize}
	cursor, err := decodeCursor(filters)
	if err != nil {
		return result, err
	}
	original := filters
	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 cursor.ID != 0 {
		clause := " AND "
		if where == "" {
			clause = " WHERE "
		}
		op := "<"
		if cursor.Before {
			op = ">"
		}
		args = append(args, cursor.ID)
		if filters.Sort == "title" {
			if cursor.Before {
				op = "<"
			} else {
				op = ">"
			}
			args = append(args, cursor.Key)
			where += fmt.Sprintf("%s(p.title_key,p.id) %s ($%d::text COLLATE \"C\",$%d)", clause, op, len(args), len(args)-1)
		} else {
			where += fmt.Sprintf("%sp.id %s $%d", clause, op, len(args))
		}
	}
	order := "p.id DESC"
	if filters.Sort == "title" {
		order = "p.title_key,p.id"
	}
	if cursor.Before {
		order = "p.id"
		if filters.Sort == "title" {
			order = "p.title_key DESC,p.id DESC"
		}
	}
	args = append(args, filters.PageSize+1)
	columns := "p.id,p.title_key"
	if filters.Sort != "title" {
		columns = "p.id, ''::text AS title_key"
	}
	// Search candidates are materialized to avoid scanning millions of wide
	// rows in ID order for sparse terms. There is no count or deep OFFSET.
	query := fmt.Sprintf("SELECT %s FROM {schema}papers p%s ORDER BY %s LIMIT $%d", columns, where, order, len(args))
	if filters.Query != "" || filters.Topic != "" {
		query = fmt.Sprintf("WITH matches AS MATERIALIZED (SELECT %s FROM {schema}papers p%s) SELECT p.id,p.title_key FROM matches p ORDER BY %s LIMIT $%d", columns, where, order, len(args))
	}
	positions := []position{}
	topicOnly := filters.Sort != "title" && filters.Topic != "" && filters.Query == ""
	if topicOnly {
		// The published paper_tags foreign key guarantees these IDs exist.
		// Walk the (tag_id,paper_id) index instead of joining all matching
		// paper rows before LIMIT. DISTINCT removes duplicate source evidence.
		topicArgs := []any{filters.Topic, filters.PageSize + 1}
		predicate := ""
		if cursor.ID != 0 {
			op := "<"
			if cursor.Before {
				op = ">"
			}
			topicArgs = append(topicArgs, cursor.ID)
			predicate = " AND pt.paper_id " + op + " $3"
		}
		direction := "DESC"
		if cursor.Before {
			direction = "ASC"
		}
		query := fmt.Sprintf(`SELECT DISTINCT pt.paper_id,''::text AS title_key FROM {schema}paper_tags pt
		 WHERE pt.tag_id=(SELECT id FROM {schema}tags WHERE namespace='arxiv:category' AND value=$1)%s
		 ORDER BY pt.paper_id %s LIMIT $2`, predicate, direction)
		positions, err = r.positions(ctx, tx, query, topicArgs)
		if err != nil {
			return result, err
		}
	}
	if !topicOnly && filters.Sort != "title" && (filters.Query != "" || filters.Topic != "") {
		// Common terms can fill a page near its cursor. Probe a bounded ID
		// interval first; never scan the entire heap in order for a rare term.
		// If the interval has fewer than size+1 matches, discard it and use
		// the complete indexed search below. No matches are silently omitted.
		anchor := cursor.ID
		if anchor == 0 {
			if err = tx.QueryRowContext(ctx, r.query(`SELECT coalesce(max(id),0) FROM {schema}papers`)).Scan(&anchor); err != nil {
				return result, err
			}
		}
		for _, window := range []int64{512, 4096} {
			bound, op := int64(0), ">="
			if anchor > window {
				bound = anchor - window
			}
			if cursor.Before {
				bound, op = int64(9223372036854775807), "<="
				if anchor <= 9223372036854775807-window {
					bound = anchor + window
				}
			}
			boundedArgs := append(append([]any{}, args...), bound)
			bounded := fmt.Sprintf("SELECT %s FROM {schema}papers p%s AND p.id %s $%d ORDER BY %s LIMIT $%d", columns, where, op, len(boundedArgs), order, len(args))
			positions, err = r.positions(ctx, tx, bounded, boundedArgs)
			if err != nil {
				return result, err
			}
			// Expand once only when there is evidence of nearby matches.
			// Zero matches goes straight to the sparse-term index fallback.
			if len(positions) == 0 || len(positions) >= filters.PageSize+1 {
				break
			}
		}
	}
	if !topicOnly && len(positions) < filters.PageSize+1 {
		positions, err = r.positions(ctx, tx, query, args)
		if err != nil {
			return result, err
		}
	}
	more := len(positions) > filters.PageSize
	if more {
		positions = positions[:filters.PageSize]
	}
	if cursor.Before {
		for i, j := 0, len(positions)-1; i < j; i, j = i+1, j-1 {
			positions[i], positions[j] = positions[j], positions[i]
		}
	}
	if len(positions) > 0 {
		if (!cursor.Before && cursor.ID != 0) || (cursor.Before && more) {
			result.PreviousCursor = encodeCursor(original, positions[0], true)
		}
		if (!cursor.Before && more) || (cursor.Before && cursor.ID != 0) {
			result.NextCursor = encodeCursor(original, positions[len(positions)-1], false)
		}
		ids := make([]int64, len(positions))
		byID := map[string]Paper{}
		for i, p := range positions {
			ids[i] = p.ID
		}
		rows, err := tx.QueryContext(ctx, r.query(`SELECT id,title,abstract,authors_json,first_published FROM {schema}papers WHERE id=ANY($1)`), ids)
		if err != nil {
			return result, err
		}
		for rows.Next() {
			p, scanErr := r.scanPaper(rows)
			if scanErr != nil {
				_ = rows.Close()
				return result, scanErr
			}
			byID[p.ID] = p
		}
		err = rows.Err()
		_ = rows.Close()
		if err != nil {
			return result, err
		}
		for _, p := range positions {
			paper, exists := byID[strconv.FormatInt(p.ID, 10)]
			if !exists {
				return result, errors.New("inconsistent paper index")
			}
			result.Items = append(result.Items, paper)
		}
	}
	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 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
}