Refine

db/db.go

@@ -0,0 +1,86 @@
+package db
+
+import (
+	"database/sql"
+	"fmt"
+	"net/url"
+	"strconv"
+	"time"
+
+	_ "github.com/jackc/pgx/v5/stdlib" // PostgreSQL driver
+)
+
+const defaultMaxOpenConns = 10
+
+// Open opens a PostgreSQL database connection pool from a DSN string. It supports custom
+// query parameters for pool configuration: pool_max_conns (default 10), pool_max_idle_conns,
+// pool_conn_max_lifetime, and pool_conn_max_idle_time. These parameters are stripped from
+// the DSN before passing it to the driver.
+func Open(dsn string) (*sql.DB, error) {
+	u, err := url.Parse(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("invalid database URL: %w", err)
+	}
+	q := u.Query()
+	maxOpenConns, err := extractIntParam(q, "pool_max_conns", defaultMaxOpenConns)
+	if err != nil {
+		return nil, err
+	}
+	maxIdleConns, err := extractIntParam(q, "pool_max_idle_conns", 0)
+	if err != nil {
+		return nil, err
+	}
+	connMaxLifetime, err := extractDurationParam(q, "pool_conn_max_lifetime", 0)
+	if err != nil {
+		return nil, err
+	}
+	connMaxIdleTime, err := extractDurationParam(q, "pool_conn_max_idle_time", 0)
+	if err != nil {
+		return nil, err
+	}
+	u.RawQuery = q.Encode()
+	db, err := sql.Open("pgx", u.String())
+	if err != nil {
+		return nil, err
+	}
+	db.SetMaxOpenConns(maxOpenConns)
+	if maxIdleConns > 0 {
+		db.SetMaxIdleConns(maxIdleConns)
+	}
+	if connMaxLifetime > 0 {
+		db.SetConnMaxLifetime(connMaxLifetime)
+	}
+	if connMaxIdleTime > 0 {
+		db.SetConnMaxIdleTime(connMaxIdleTime)
+	}
+	if err := db.Ping(); err != nil {
+		return nil, fmt.Errorf("ping failed: %w", err)
+	}
+	return db, nil
+}
+
+func extractIntParam(q url.Values, key string, defaultValue int) (int, error) {
+	s := q.Get(key)
+	if s == "" {
+		return defaultValue, nil
+	}
+	q.Del(key)
+	v, err := strconv.Atoi(s)
+	if err != nil {
+		return 0, fmt.Errorf("invalid %s value %q: %w", key, s, err)
+	}
+	return v, nil
+}
+
+func extractDurationParam(q url.Values, key string, defaultValue time.Duration) (time.Duration, error) {
+	s := q.Get(key)
+	if s == "" {
+		return defaultValue, nil
+	}
+	q.Del(key)
+	d, err := time.ParseDuration(s)
+	if err != nil {
+		return 0, fmt.Errorf("invalid %s value %q: %w", key, s, err)
+	}
+	return d, nil
+}

db/test/test.go

@@ -0,0 +1,63 @@
+package dbtest
+
+import (
+	"database/sql"
+	"fmt"
+	"net/url"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"heckel.io/ntfy/v2/db"
+	"heckel.io/ntfy/v2/util"
+)
+
+const testPoolMaxConns = "2"
+
+// CreateTestSchema creates a temporary PostgreSQL schema and returns the DSN pointing to it.
+// It registers a cleanup function to drop the schema when the test finishes.
+// If NTFY_TEST_DATABASE_URL is not set, the test is skipped.
+func CreateTestSchema(t *testing.T) string {
+	t.Helper()
+	dsn := os.Getenv("NTFY_TEST_DATABASE_URL")
+	if dsn == "" {
+		t.Skip("NTFY_TEST_DATABASE_URL not set")
+	}
+	schema := fmt.Sprintf("test_%s", util.RandomString(10))
+	u, err := url.Parse(dsn)
+	require.Nil(t, err)
+	q := u.Query()
+	q.Set("pool_max_conns", testPoolMaxConns)
+	u.RawQuery = q.Encode()
+	dsn = u.String()
+	setupDB, err := db.Open(dsn)
+	require.Nil(t, err)
+	_, err = setupDB.Exec(fmt.Sprintf("CREATE SCHEMA %s", schema))
+	require.Nil(t, err)
+	require.Nil(t, setupDB.Close())
+	q.Set("search_path", schema)
+	u.RawQuery = q.Encode()
+	schemaDSN := u.String()
+	t.Cleanup(func() {
+		cleanDB, err := db.Open(dsn)
+		if err == nil {
+			cleanDB.Exec(fmt.Sprintf("DROP SCHEMA %s CASCADE", schema))
+			cleanDB.Close()
+		}
+	})
+	return schemaDSN
+}
+
+// CreateTestDB creates a temporary PostgreSQL schema and returns an open *sql.DB connection to it.
+// It registers cleanup functions to close the DB and drop the schema when the test finishes.
+// If NTFY_TEST_DATABASE_URL is not set, the test is skipped.
+func CreateTestDB(t *testing.T) *sql.DB {
+	t.Helper()
+	schemaDSN := CreateTestSchema(t)
+	testDB, err := db.Open(schemaDSN)
+	require.Nil(t, err)
+	t.Cleanup(func() {
+		testDB.Close()
+	})
+	return testDB
+}