Use one PG connection, add support for connection params

This commit is contained in:
binwiederhier
2026-02-19 22:34:53 -05:00
parent 305e3fc9af
commit 4cbd80c68e
12 changed files with 172 additions and 95 deletions

View File

@@ -25,7 +25,6 @@ var errNoRows = errors.New("no rows found")
type Store interface {
AddMessage(m *model.Message) error
AddMessages(ms []*model.Message) error
DB() *sql.DB
Message(id string) (*model.Message, error)
MessageCounts() (map[string]int, error)
Messages(topic string, since model.SinceMarker, scheduled bool) ([]*model.Message, error)
@@ -99,11 +98,6 @@ func newCommonStore(db *sql.DB, queries storeQueries, batchSize int, batchTimeou
return c
}
// DB returns the underlying database connection
func (c *commonStore) DB() *sql.DB {
return c.db
}
// AddMessage stores a message to the message cache synchronously, or queues it to be stored at a later date asynchronously.
func (c *commonStore) AddMessage(m *model.Message) error {
if c.queue != nil {

View File

@@ -3,8 +3,6 @@ package message
import (
"database/sql"
"time"
_ "github.com/jackc/pgx/v5/stdlib" // PostgreSQL driver
)
// PostgreSQL runtime query constants
@@ -104,16 +102,8 @@ var pgQueries = storeQueries{
updateMessageTime: pgUpdateMessageTimesQuery,
}
// NewPostgresStore creates a new PostgreSQL-backed message cache store.
func NewPostgresStore(dsn string, batchSize int, batchTimeout time.Duration) (Store, error) {
db, err := sql.Open("pgx", dsn)
if err != nil {
return nil, err
}
db.SetMaxOpenConns(25)
if err := db.Ping(); err != nil {
return nil, err
}
// NewPostgresStore creates a new PostgreSQL-backed message cache store using an existing database connection pool.
func NewPostgresStore(db *sql.DB, batchSize int, batchTimeout time.Duration) (Store, error) {
if err := setupPostgresDB(db); err != nil {
return nil, err
}

View File

@@ -1,7 +1,6 @@
package message_test
import (
"database/sql"
"fmt"
"net/url"
"os"
@@ -9,6 +8,7 @@ import (
"github.com/stretchr/testify/require"
"heckel.io/ntfy/v2/message"
"heckel.io/ntfy/v2/postgres"
"heckel.io/ntfy/v2/util"
)
@@ -17,24 +17,25 @@ func newTestPostgresStore(t *testing.T) message.Store {
if dsn == "" {
t.Skip("NTFY_TEST_DATABASE_URL not set, skipping PostgreSQL tests")
}
// Create a unique schema for this test
schema := fmt.Sprintf("test_%s", util.RandomString(10))
setupDB, err := sql.Open("pgx", dsn)
require.Nil(t, err)
_, err = setupDB.Exec(fmt.Sprintf("CREATE SCHEMA %s", schema))
require.Nil(t, err)
require.Nil(t, setupDB.Close())
// Open store with search_path set to the new schema
u, err := url.Parse(dsn)
require.Nil(t, err)
q := u.Query()
q.Set("search_path", schema)
u.RawQuery = q.Encode()
store, err := message.NewPostgresStore(u.String(), 0, 0)
schemaDSN := u.String()
setupDB, err := postgres.OpenDB(dsn)
require.Nil(t, err)
_, err = setupDB.Exec(fmt.Sprintf("CREATE SCHEMA %s", schema))
require.Nil(t, err)
require.Nil(t, setupDB.Close())
db, err := postgres.OpenDB(schemaDSN)
require.Nil(t, err)
store, err := message.NewPostgresStore(db, 0, 0)
require.Nil(t, err)
t.Cleanup(func() {
store.Close()
cleanDB, err := sql.Open("pgx", dsn)
cleanDB, err := postgres.OpenDB(dsn)
if err == nil {
cleanDB.Exec(fmt.Sprintf("DROP SCHEMA %s CASCADE", schema))
cleanDB.Close()