ntfy/message/cache_postgres.go

package message

import (
	"database/sql"
	"strconv"
	"strings"
	"time"

	"heckel.io/ntfy/v2/db"
	"heckel.io/ntfy/v2/db/schema"
)

// PostgreSQL runtime query constants
const (
	postgresSelectScheduledMessageIDsBySeqIDQuery = `SELECT mid FROM message WHERE topic = $1 AND sequence_id = $2 AND published = FALSE`
	postgresDeleteScheduledBySequenceIDQuery      = `DELETE FROM message WHERE topic = $1 AND sequence_id = $2 AND published = FALSE`
	postgresUpdateMessagesForTopicExpiryQuery     = `UPDATE message SET expires = $1 WHERE topic = $2`
	postgresSelectMessagesByIDQuery               = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE mid = $1
	`
	postgresSelectMessagesSinceTimeQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE topic = $1 AND time >= $2 AND published = TRUE
		ORDER BY time DESC, id DESC
	`
	postgresSelectMessagesSinceTimeIncludeScheduledQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE topic = $1 AND time >= $2
		ORDER BY time DESC, id DESC
	`
	postgresSelectMessageRowIDQuery    = `SELECT id FROM message WHERE mid = $1 LIMIT 1`
	postgresSelectMessagesSinceIDQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE topic = $1
		  AND id > $2
		  AND published = TRUE
		ORDER BY time DESC, id DESC
	`
	postgresSelectMessagesSinceIDIncludeScheduledQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE topic = $1
		  AND (id > $2 OR published = FALSE)
		ORDER BY time DESC, id DESC
	`
	postgresSelectMessagesLatestQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE topic = $1 AND published = TRUE
		ORDER BY time DESC, id DESC
		LIMIT 1
	`
	postgresSelectMessagesDueQuery = `
		SELECT mid, sequence_id, time, event, expires, topic, message, title, priority, tags, click, icon, actions, attachment_name, attachment_type, attachment_size, attachment_expires, attachment_url, sender, user_id, content_type, encoding
		FROM message
		WHERE time <= $1 AND published = FALSE
		ORDER BY time, id
	`
	postgresUpdateMessagePublishedQuery = `UPDATE message SET published = TRUE WHERE mid = $1`
	// Planner estimate, since a COUNT(*) scans the whole table; reltuples is -1 if never analyzed
	postgresSelectMessagesCountQuery = `
		SELECT CASE WHEN reltuples < 0 THEN (SELECT COUNT(*) FROM message) ELSE reltuples::BIGINT END
		FROM pg_class
		WHERE oid = 'message'::regclass
	`
	postgresSelectTopicsQuery = `SELECT topic FROM message GROUP BY topic`

	postgresDeleteExpiredMessagesQuery         = `DELETE FROM message WHERE mid IN (SELECT mid FROM message WHERE expires <= $1 AND published = TRUE LIMIT $2)`
	postgresMarkExpiredAttachmentsDeletedQuery = `UPDATE message SET attachment_deleted = TRUE WHERE mid IN (SELECT mid FROM message WHERE attachment_expires > 0 AND attachment_expires <= $1 AND attachment_deleted = FALSE LIMIT $2)`
	postgresSelectAttachmentsSizeBySenderQuery = `SELECT COALESCE(SUM(attachment_size), 0) FROM message WHERE user_id = '' AND sender = $1 AND attachment_expires >= $2`
	postgresSelectAttachmentsSizeByUserIDQuery = `SELECT COALESCE(SUM(attachment_size), 0) FROM message WHERE user_id = $1 AND attachment_expires >= $2`
	postgresSelectAttachmentsWithSizesQuery    = `SELECT mid, attachment_size FROM message WHERE attachment_expires > $1 AND attachment_deleted = FALSE`

	postgresSelectStatsQuery       = `SELECT value FROM message_stats WHERE key = 'messages'`
	postgresUpdateStatsQuery       = `UPDATE message_stats SET value = $1 WHERE key = 'messages'`
	postgresUpdateMessageTimeQuery = `UPDATE message SET time = $1 WHERE mid = $2`
)

var postgresQueries = queries{
	insertMessages:                   postgresInsertMessages,
	selectScheduledMessageIDsBySeqID: postgresSelectScheduledMessageIDsBySeqIDQuery,
	deleteScheduledBySequenceID:      postgresDeleteScheduledBySequenceIDQuery,
	updateMessagesForTopicExpiry:     postgresUpdateMessagesForTopicExpiryQuery,
	selectMessagesByID:               postgresSelectMessagesByIDQuery,
	selectMessagesSinceTime:          postgresSelectMessagesSinceTimeQuery,
	selectMessagesSinceTimeScheduled: postgresSelectMessagesSinceTimeIncludeScheduledQuery,
	selectMessageRowID:               postgresSelectMessageRowIDQuery,
	selectMessagesSinceID:            postgresSelectMessagesSinceIDQuery,
	selectMessagesSinceIDScheduled:   postgresSelectMessagesSinceIDIncludeScheduledQuery,
	selectMessagesLatest:             postgresSelectMessagesLatestQuery,
	selectMessagesDue:                postgresSelectMessagesDueQuery,
	deleteExpiredMessages:            postgresDeleteExpiredMessagesQuery,
	updateMessagePublished:           postgresUpdateMessagePublishedQuery,
	selectMessagesCount:              postgresSelectMessagesCountQuery,
	selectTopics:                     postgresSelectTopicsQuery,
	markExpiredAttachmentsDeleted:    postgresMarkExpiredAttachmentsDeletedQuery,
	selectAttachmentsSizeBySender:    postgresSelectAttachmentsSizeBySenderQuery,
	selectAttachmentsSizeByUserID:    postgresSelectAttachmentsSizeByUserIDQuery,
	selectAttachmentsWithSizes:       postgresSelectAttachmentsWithSizesQuery,
	selectStats:                      postgresSelectStatsQuery,
	updateStats:                      postgresUpdateStatsQuery,
	updateMessageTime:                postgresUpdateMessageTimeQuery,
}

// NewPostgresStore creates a new PostgreSQL-backed message cache store using an existing database connection pool.
func NewPostgresStore(d *db.DB, batchSize int, batchTimeout time.Duration) (*Cache, error) {
	if err := schema.Migrate(d.Primary(), schema.Postgres, schemaStore, postgresCurrentSchemaVersion, postgresCreateTables, postgresMigrations); err != nil {
		return nil, err
	}
	return newCache(d, postgresQueries, nil, batchSize, batchTimeout, false), nil
}

// postgresInsertMessages writes all rows in one statement, transposed into one array parameter
// per column: INSERT INTO message (c1, ...) SELECT * FROM unnest($1::T1[], ...). The statement
// text is the same for every batch size: pgx prepares each distinct text server-side and keeps up
// to 512 per connection, and a multi-row INSERT per size left hundreds of multi-MB plans in every
// pooled backend. Building the text per batch is a few string joins, nothing next to the round trip.
func postgresInsertMessages(tx *sql.Tx, rows [][]insertValue) error {
	columns, arrays, args := make([]string, 0, len(rows[0])), make([]string, 0, len(rows[0])), make([]any, 0, len(rows[0]))
	for i, v := range rows[0] {
		column := make([]any, 0, len(rows))
		for _, row := range rows {
			column = append(column, row[i].value)
		}
		columns = append(columns, v.column)
		arrays = append(arrays, "$"+strconv.Itoa(i+1)+"::"+v.sqlType+"[]")
		args = append(args, column)
	}
	query := "INSERT INTO message (" + strings.Join(columns, ", ") + ") SELECT * FROM unnest(" + strings.Join(arrays, ", ") + ")"
	_, err := tx.Exec(query, args...)
	return err
}