Warp/server/internal/storage/postgres/postgres.go

// Package postgres implements the storage interfaces on PostgreSQL 15+ via
// pgx's database/sql driver. This is the production backend and the one
// Complement runs against.
package postgres

import (
	"context"
	"database/sql"
	"embed"
	"fmt"
	"log/slog"

	_ "github.com/jackc/pgx/v5/stdlib" // registers the "pgx" database/sql driver

	"github.com/dariuszjaniak/warp/server/internal/storage"
	"github.com/dariuszjaniak/warp/server/internal/storage/migrations"
	"github.com/dariuszjaniak/warp/server/internal/storage/sqlstore"
)

//go:embed migrations/*.sql
var migrationFiles embed.FS

func init() {
	storage.Register(storage.DriverPostgres, Open)
}

var dialect = migrations.Dialect{
	Name: storage.DriverPostgres,
	CreateTable: `CREATE TABLE IF NOT EXISTS schema_migrations (
		version    INTEGER     PRIMARY KEY,
		name       TEXT        NOT NULL,
		checksum   TEXT        NOT NULL,
		applied_at TIMESTAMPTZ NOT NULL
	)`,
	SelectApplied: `SELECT version, name, checksum, applied_at FROM schema_migrations ORDER BY version`,
	InsertApplied: `INSERT INTO schema_migrations (version, name, checksum, applied_at) VALUES ($1, $2, $3, $4)`,
}

var settingsQueries = sqlstore.SettingsQueries{
	Get:    `SELECT setting_value FROM server_settings WHERE setting_key = $1`,
	GetAll: `SELECT setting_key, setting_value FROM server_settings`,
	Upsert: `INSERT INTO server_settings (setting_key, setting_value, updated_at) VALUES ($1, $2, $3)
		ON CONFLICT (setting_key) DO UPDATE SET setting_value = excluded.setting_value, updated_at = excluded.updated_at`,
	Delete: `DELETE FROM server_settings WHERE setting_key = $1`,
}

// Database is the PostgreSQL implementation of storage.Database.
type Database struct {
	*sqlstore.Store
}

// Open connects to the DSN in cfg.ConnectionString, e.g.
// "postgres://warp:secret@localhost/warp?sslmode=disable".
func Open(ctx context.Context, cfg storage.Config, log *slog.Logger) (storage.Database, error) {
	if cfg.ConnectionString == "" {
		return nil, fmt.Errorf("postgres: connection_string must not be empty")
	}
	db, err := sql.Open("pgx", cfg.ConnectionString)
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(cfg.MaxOpenConns)
	db.SetMaxIdleConns(cfg.MaxIdleConns)
	db.SetConnMaxLifetime(cfg.ConnMaxLifetime)

	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("postgres: connect: %w", err)
	}

	store, err := sqlstore.New(sqlstore.Options{
		DriverName:        storage.DriverPostgres,
		DB:                db,
		Dialect:           dialect,
		MigrationsFS:      migrationFiles,
		MigrationsDir:     "migrations",
		Settings:          settingsQueries,
		Rebind:            rebind,
		IsUniqueViolation: isUniqueViolation,
		Logger:            log,
	})
	if err != nil {
		_ = db.Close()
		return nil, err
	}
	return &Database{Store: store}, nil
}