1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
package main

// Database setup and schema migration.
//
// The schema is ported verbatim from the Haskell `migrate` in
// src/WhatcdResolver.hs. It is deliberately *not* modernised: the seeding
// weight is computed by a plpgsql function inside a STORED generated column, so
// changing the function or the JSON it reads silently changes the ranking of
// every torrent in the database. Any change here must be a new migration.
//
// Unlike the Haskell version, which started its own PostgreSQL via tmp-postgres,
// this connects to an already-running server (see whatcd-resolver.1, section
// DATABASE).

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"go.opentelemetry.io/otel/trace"
)

// migrationSQL is the whole schema, idempotent, applied at every startup.
//
// Copied verbatim from the Haskell version, including the comments, which
// document intent that is not otherwise recoverable (in particular the reasons
// behind each factor of calc_seeding_weight).
const migrationSQL = `
    CREATE SCHEMA IF NOT EXISTS redacted;

    CREATE TABLE IF NOT EXISTS redacted.settings (
      id SERIAL PRIMARY KEY,
      key TEXT NOT NULL UNIQUE,
      value JSONB
    );

    CREATE TABLE IF NOT EXISTS redacted.torrent_groups (
      id SERIAL PRIMARY KEY,
      group_id INTEGER,
      group_name TEXT,
      full_json_result JSONB,
      UNIQUE(group_id)
    );

    CREATE TABLE IF NOT EXISTS redacted.torrents_json (
      id SERIAL PRIMARY KEY,
      torrent_id INTEGER,
      torrent_group SERIAL NOT NULL REFERENCES redacted.torrent_groups(id) ON DELETE CASCADE,
      full_json_result JSONB,
      UNIQUE(torrent_id)
    );

    CREATE INDEX IF NOT EXISTS redacted_torrents_json_torrent_group_fk ON redacted.torrents_json (torrent_group);


    ALTER TABLE redacted.torrents_json
    ADD COLUMN IF NOT EXISTS torrent_file bytea NULL;
    ALTER TABLE redacted.torrents_json
    ADD COLUMN IF NOT EXISTS transmission_torrent_hash text NULL;

    -- the seeding weight is used to find the best torrent in a group.
    CREATE OR REPLACE FUNCTION calc_seeding_weight(full_json_result jsonb) RETURNS int AS $$
    BEGIN
      RETURN
        -- three times seeders plus one times snatches
        (3 * (full_json_result->'seeders')::integer
        + (full_json_result->'snatches')::integer
        )
        -- prefer remasters by multiplying them with 3
        * (CASE
            WHEN full_json_result->>'remasterTitle' ILIKE '%remaster%'
            THEN 3
            ELSE 1
          END)
        -- slightly push mp3 V0, to make sure it’s preferred over 320 CBR
        * (CASE
            WHEN full_json_result->>'encoding' ILIKE '%v0%'
            THEN 2
            ELSE 1
          END)
        -- remove 24bit torrents from the result (wayyy too big)
        * (CASE
            WHEN full_json_result->>'encoding' ILIKE '%24bit%'
            THEN 0
            ELSE 1
          END)
        -- discount FLACS, so we only use them when there’s no mp3 alternative (to save space)
        / (CASE
            WHEN full_json_result->>'encoding' ILIKE '%lossless%'
            THEN 5
            ELSE 1
          END)
        ;
    END;
    $$ LANGUAGE plpgsql IMMUTABLE;

    ALTER TABLE redacted.torrents_json
    ADD COLUMN IF NOT EXISTS seeding_weight int NOT NULL GENERATED ALWAYS AS (calc_seeding_weight(full_json_result)) STORED;

    CREATE OR REPLACE FUNCTION artist_record_to_id(artists jsonb) RETURNS int[]
    as $$
      SELECT array_agg(x::int) from jsonb_path_query(artists, '$[*].id') j(x);
    $$ LANGUAGE sql IMMUTABLE;

    ALTER TABLE redacted.torrents_json
    ADD COLUMN IF NOT EXISTS artist_ids int[] NOT NULL GENERATED ALWAYS AS (COALESCE(artist_record_to_id(full_json_result->'artists'), ARRAY[]::int[])) STORED;

    CREATE INDEX IF NOT EXISTS torrents_json_artist_ids ON redacted.torrents_json USING GIN (artist_ids);

    -- inflect out values of the full json.
    CREATE OR REPLACE VIEW redacted.torrents AS
    SELECT
      t.id,
      t.torrent_id,
      t.torrent_group,
      -- the seeding weight is used to find the best torrent in a group.
      t.seeding_weight,
      t.full_json_result,
      t.torrent_file,
      t.transmission_torrent_hash,
      t.artist_ids
    FROM redacted.torrents_json t;

    CREATE INDEX IF NOT EXISTS torrents_json_seeding ON redacted.torrents_json(((full_json_result->'seeding')::integer));
    CREATE INDEX IF NOT EXISTS torrents_json_snatches ON redacted.torrents_json(((full_json_result->'snatches')::integer));

    CREATE TABLE IF NOT EXISTS redacted.artist_favourites (
      id SERIAL PRIMARY KEY,
      artist_id INTEGER NOT NULL,
      UNIQUE(artist_id)
    );

    -- table for storing related/similar artists from the API
    CREATE TABLE IF NOT EXISTS redacted.similar_artists (
      id SERIAL PRIMARY KEY,
      artist_id INTEGER NOT NULL,
      similar_artist_id INTEGER NOT NULL,
      similar_artist_name TEXT NOT NULL,
      score INTEGER NOT NULL,
      created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
      UNIQUE(artist_id, similar_artist_id)
    );

    CREATE INDEX IF NOT EXISTS similar_artists_artist_id_idx ON redacted.similar_artists (artist_id);
    CREATE INDEX IF NOT EXISTS similar_artists_similar_artist_id_idx ON redacted.similar_artists (similar_artist_id);

    -- fast lookup table for artist id -> name mapping
    CREATE TABLE IF NOT EXISTS redacted.artists (
      artist_id INTEGER PRIMARY KEY,
      artist_name TEXT NOT NULL,
      updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
    );
`

// connectDB opens the pool and applies the migration.
func connectDB(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, fmt.Errorf("parsing database url: %w", err)
	}
	// The Haskell version kept a 20-connection pool with a 600s idle timeout.
	cfg.MaxConns = 20
	cfg.MaxConnIdleTime = 600 * time.Second

	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, fmt.Errorf("connecting to database: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("pinging database: %w", err)
	}
	return pool, nil
}

// migrate applies the schema. Safe to run repeatedly.
func migrate(ctx context.Context, pool *pgxpool.Pool) error {
	return inSpan(ctx, "Database Migration", func(ctx context.Context, span trace.Span) error {
		if _, err := pool.Exec(ctx, migrationSQL); err != nil {
			return fmt.Errorf("running migration: %w", err)
		}
		slog.Info("database migration applied")
		return nil
	})
}

// ---------------------------------------------------------------------------
// Query helpers
//
// pgx already returns good errors, so these only add tracing and the
// "exactly one row" checks the Haskell code made explicit.
// ---------------------------------------------------------------------------

// txFunc runs f inside a transaction, committing on success and rolling back on
// error or panic. Replaces the Haskell `runTransaction`.
func withTx(ctx context.Context, pool *pgxpool.Pool, f func(context.Context, pgx.Tx) error) error {
	tx, err := pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("beginning transaction: %w", err)
	}
	defer func() {
		// Rollback after a successful commit is a no-op, so this is safe
		// unconditionally and also covers panics.
		_ = tx.Rollback(ctx)
	}()
	if err := f(ctx, tx); err != nil {
		return err
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("committing transaction: %w", err)
	}
	return nil
}

// assertOneUpdated mirrors the Haskell helper of the same name: several updates
// address a single row by primary key and a different count means the
// assumption behind the query no longer holds.
func assertOneUpdated(name string, affected int64) error {
	if affected != 1 {
		return fmt.Errorf("%s: expected to update exactly one row, but updated %d row(s)", name, affected)
	}
	return nil
}