diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 29cf2c0..193b2ce 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -15,6 +15,21 @@ jobs: platform: [ubuntu-latest] runs-on: ubuntu-latest + services: + postgres: + image: postgres:16 + env: + POSTGRES_USER: gorm + POSTGRES_PASSWORD: gorm + POSTGRES_DB: gorm + ports: + - 9920:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + steps: - name: Set up Go 1.x uses: actions/setup-go@v5 @@ -29,4 +44,6 @@ jobs: run: go build . - name: Run tests + env: + POSTGRES_DSN: host=localhost user=gorm password=gorm dbname=gorm port=9920 sslmode=disable run: go test -race -count=1 -v ./... diff --git a/example_listener_test.go b/example_listener_test.go new file mode 100644 index 0000000..113614e --- /dev/null +++ b/example_listener_test.go @@ -0,0 +1,48 @@ +package postgres_test + +import ( + "context" + "fmt" + "log" + + "gorm.io/driver/postgres" + "gorm.io/gorm" +) + +// ExampleNewListener demonstrates receiving a PostgreSQL notification. The +// listener reserves one dedicated connection until Close is called, and its +// LISTEN registrations last only for that connection's session. +func ExampleNewListener() { + dsn := "host=localhost user=gorm password=gorm dbname=gorm port=9920 sslmode=disable" + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) + if err != nil { + log.Fatal(err) + } + + ctx := context.Background() + + listener, err := postgres.NewListener(ctx, db) + if err != nil { + log.Fatal(err) + } + defer listener.Close() + + if err := listener.Listen(ctx, "events"); err != nil { + log.Fatal(err) + } + + // Notifications can be published through GORM. Inside a transaction, + // PostgreSQL delivers them only after the transaction commits. + if err := db.WithContext(ctx).Exec("SELECT pg_notify(?, ?)", "events", "hello").Error; err != nil { + log.Fatal(err) + } + + // WaitForNotification blocks until a notification arrives or ctx is + // done; cancel the context to stop waiting. + notification, err := listener.WaitForNotification(ctx) + if err != nil { + log.Fatal(err) + } + + fmt.Println(notification.Channel, notification.Payload) +} diff --git a/listener.go b/listener.go new file mode 100644 index 0000000..19db61d --- /dev/null +++ b/listener.go @@ -0,0 +1,234 @@ +package postgres + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "strings" + "sync/atomic" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/stdlib" + "gorm.io/gorm" +) + +var ( + // ErrUnsupportedConnection is returned by NewListener when the underlying + // database connection is not managed by the pgx stdlib driver, for example + // when a custom Config.Conn or Config.DriverName routes connections + // through a different database/sql driver. + ErrUnsupportedConnection = errors.New("postgres: LISTEN requires a pgx stdlib connection") + + // ErrListenerClosed is returned by all Listener methods after Close has + // been called. + ErrListenerClosed = errors.New("postgres: listener is closed") + + // ErrNotificationHandlerConflict is returned by WaitForNotification when + // the pgx connection has a custom OnNotification handler. Such a handler + // consumes notifications before pgx can return them to the Listener. + ErrNotificationHandlerConflict = errors.New("postgres: listener conflicts with configured pgx OnNotification handler") +) + +// Notification is a PostgreSQL notification received by a Listener. +type Notification struct { + // PID is the backend process ID of the PostgreSQL session that sent the + // notification. + PID uint32 + // Channel is the name of the channel the notification was sent on. + Channel string + // Payload is the optional payload string sent with the notification. + Payload string +} + +// Listener receives PostgreSQL NOTIFY messages using LISTEN. +// +// A Listener reserves one dedicated connection from the *sql.DB behind the +// *gorm.DB for its entire lifetime, so that connection is unavailable to +// other queries until Close is called. A pool limited to a single connection +// (for example via SetMaxOpenConns(1)) cannot run other queries, such as +// sending notifications, while a Listener exists. +// +// LISTEN registrations are scoped to the reserved PostgreSQL session: they +// disappear when the Listener is closed or the connection is lost, and +// notifications sent while no session is listening are dropped. PostgreSQL +// notifications are not a persistent queue; a Listener may miss notifications +// while it is disconnected. Reconnecting and re-subscribing after a +// connection loss is the application's responsibility - the Listener never +// reconnects on its own. +// +// Notifications can be published through GORM with +// +// db.Exec("SELECT pg_notify(?, ?)", channel, payload) +// +// When pg_notify (or NOTIFY) runs inside a transaction, PostgreSQL delivers +// the notification only after the transaction commits, and never if it rolls +// back. +// +// A Listener is not safe for concurrent use. Calls must be serialized by the +// caller, and Listen, Unlisten, UnlistenAll and WaitForNotification must not +// run concurrently with each other or with Close. To interrupt a blocked +// WaitForNotification, cancel the context passed to it; afterwards the +// Listener remains usable and may be closed or reused. +type Listener struct { + conn *sql.Conn + closed atomic.Bool +} + +// NewListener reserves a dedicated connection from the connection pool behind +// db and returns a Listener bound to that connection's PostgreSQL session. +// +// db may be a root, session or transaction *gorm.DB handle; in every case the +// Listener acquires its own independent connection from the pool and never +// joins an ongoing transaction. The connection pool must be backed by the pgx +// stdlib driver (the default for this dialector); otherwise NewListener +// returns an error wrapping ErrUnsupportedConnection. +// +// The caller must call Close to release the reserved connection. +func NewListener(ctx context.Context, db *gorm.DB) (*Listener, error) { + if db == nil { + return nil, gorm.ErrInvalidDB + } + + sqlDB, err := db.DB() + if err != nil { + return nil, fmt.Errorf("postgres: resolve *sql.DB for listener: %w", err) + } + + conn, err := sqlDB.Conn(ctx) + if err != nil { + return nil, fmt.Errorf("postgres: acquire listener connection: %w", err) + } + + supported := false + if err := conn.Raw(func(driverConn any) error { + _, supported = driverConn.(*stdlib.Conn) + return nil + }); err != nil { + _ = conn.Close() + return nil, fmt.Errorf("postgres: inspect listener connection: %w", err) + } + if !supported { + _ = conn.Close() + return nil, ErrUnsupportedConnection + } + + return &Listener{conn: conn}, nil +} + +// Listen registers the Listener's session as a listener on channel, which is +// quoted as a SQL identifier and therefore matched case-sensitively. +func (l *Listener) Listen(ctx context.Context, channel string) error { + if l.closed.Load() { + return ErrListenerClosed + } + quoted, err := quoteChannel(channel) + if err != nil { + return err + } + return l.exec(ctx, "LISTEN "+quoted) +} + +// Unlisten removes the session's registration on channel. Unlistening a +// channel that is not registered is not an error. +func (l *Listener) Unlisten(ctx context.Context, channel string) error { + if l.closed.Load() { + return ErrListenerClosed + } + quoted, err := quoteChannel(channel) + if err != nil { + return err + } + return l.exec(ctx, "UNLISTEN "+quoted) +} + +// UnlistenAll removes all of the session's channel registrations using +// UNLISTEN *. +func (l *Listener) UnlistenAll(ctx context.Context) error { + return l.exec(ctx, "UNLISTEN *") +} + +// WaitForNotification blocks until a notification is received on one of the +// registered channels or ctx is done. +// +// When ctx is canceled or times out, the returned error satisfies +// errors.Is(err, context.Canceled) or errors.Is(err, context.DeadlineExceeded) +// and the Listener remains usable. Any other error usually means the +// underlying connection is broken; the Listener should then be closed and, if +// desired, replaced with a new one. If the pgx connection was configured with +// a custom OnNotification handler, that handler consumes the notification and +// WaitForNotification returns ErrNotificationHandlerConflict. +func (l *Listener) WaitForNotification(ctx context.Context) (*Notification, error) { + var notification *Notification + err := l.raw(func(conn *pgx.Conn) error { + n, err := conn.WaitForNotification(ctx) + if err != nil { + return err + } + if n == nil { + return ErrNotificationHandlerConflict + } + notification = &Notification{PID: n.PID, Channel: n.Channel, Payload: n.Payload} + return nil + }) + if err != nil { + return nil, err + } + return notification, nil +} + +// Close releases the reserved connection. Because the session may still hold +// LISTEN registrations and buffered notifications, the physical connection is +// discarded instead of being reused; its slot in the pool is freed either +// way. Close is idempotent: the first call releases the connection and +// subsequent calls return nil. All other methods return ErrListenerClosed +// after Close. +func (l *Listener) Close() error { + if !l.closed.CompareAndSwap(false, true) { + return nil + } + // Returning driver.ErrBadConn makes database/sql discard the underlying + // driver connection rather than returning the session to the pool. + _ = l.conn.Raw(func(any) error { return driver.ErrBadConn }) + if err := l.conn.Close(); err != nil && !errors.Is(err, sql.ErrConnDone) { + return fmt.Errorf("postgres: close listener connection: %w", err) + } + return nil +} + +// raw runs f against the reserved pgx connection. +func (l *Listener) raw(f func(conn *pgx.Conn) error) error { + if l.closed.Load() { + return ErrListenerClosed + } + return l.conn.Raw(func(driverConn any) error { + stdConn, ok := driverConn.(*stdlib.Conn) + if !ok { + return ErrUnsupportedConnection + } + return f(stdConn.Conn()) + }) +} + +func (l *Listener) exec(ctx context.Context, sql string) error { + return l.raw(func(conn *pgx.Conn) error { + if _, err := conn.Exec(ctx, sql); err != nil { + return fmt.Errorf("postgres: %s: %w", sql, err) + } + return nil + }) +} + +// quoteChannel quotes a notification channel name as a SQL identifier. +// Channel names are identifiers, not bind values, so they cannot be passed as +// ordinary query parameters. +func quoteChannel(channel string) (string, error) { + if channel == "" { + return "", errors.New("postgres: notification channel name must not be empty") + } + if strings.ContainsRune(channel, '\x00') { + return "", errors.New("postgres: notification channel name must not contain NUL") + } + return pgx.Identifier{channel}.Sanitize(), nil +} diff --git a/listener_integration_test.go b/listener_integration_test.go new file mode 100644 index 0000000..9659184 --- /dev/null +++ b/listener_integration_test.go @@ -0,0 +1,479 @@ +package postgres_test + +import ( + "context" + "errors" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "github.com/jackc/pgx/v5/stdlib" + "gorm.io/driver/postgres" + "gorm.io/gorm" +) + +// integrationDB opens a *gorm.DB against the PostgreSQL server configured via +// the POSTGRES_DSN environment variable, skipping the test when it is unset. +func integrationDB(t *testing.T) *gorm.DB { + t.Helper() + dsn := os.Getenv("POSTGRES_DSN") + if dsn == "" { + t.Skip("POSTGRES_DSN is not set, skipping integration test") + } + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) + if err != nil { + t.Fatalf("gorm.Open() error = %v", err) + } + t.Cleanup(func() { + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func newTestListener(t *testing.T, ctx context.Context, db *gorm.DB) *postgres.Listener { + t.Helper() + listener, err := postgres.NewListener(ctx, db) + if err != nil { + t.Fatalf("NewListener() error = %v", err) + } + t.Cleanup(func() { _ = listener.Close() }) + return listener +} + +func notify(t *testing.T, ctx context.Context, db *gorm.DB, channel, payload string) { + t.Helper() + if err := db.WithContext(ctx).Exec("SELECT pg_notify(?, ?)", channel, payload).Error; err != nil { + t.Fatalf("pg_notify(%q, %q) error = %v", channel, payload, err) + } +} + +func waitForNotification(t *testing.T, listener *postgres.Listener) *postgres.Notification { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + n, err := listener.WaitForNotification(ctx) + if err != nil { + t.Fatalf("WaitForNotification() error = %v", err) + } + return n +} + +func TestListener_receiveNotification(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + if err := listener.Listen(ctx, "gorm_test_events"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + + // Send the notification from a dedicated connection so the notifying + // backend's PID is known. + sqlDB, err := db.DB() + if err != nil { + t.Fatalf("db.DB() error = %v", err) + } + senderConn, err := sqlDB.Conn(ctx) + if err != nil { + t.Fatalf("sqlDB.Conn() error = %v", err) + } + defer senderConn.Close() + var senderPID uint32 + if err := senderConn.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&senderPID); err != nil { + t.Fatalf("pg_backend_pid() error = %v", err) + } + if _, err := senderConn.ExecContext(ctx, "SELECT pg_notify($1, $2)", "gorm_test_events", "hello"); err != nil { + t.Fatalf("pg_notify error = %v", err) + } + + n := waitForNotification(t, listener) + if n.Channel != "gorm_test_events" { + t.Errorf("Channel = %q, want %q", n.Channel, "gorm_test_events") + } + if n.Payload != "hello" { + t.Errorf("Payload = %q, want %q", n.Payload, "hello") + } + if n.PID != senderPID { + t.Errorf("PID = %d, want sender backend PID %d", n.PID, senderPID) + } +} + +func TestListener_multipleChannels(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + for _, channel := range []string{"gorm_test_a", "gorm_test_b"} { + if err := listener.Listen(ctx, channel); err != nil { + t.Fatalf("Listen(%q) error = %v", channel, err) + } + } + + notify(t, ctx, db, "gorm_test_a", "1") + notify(t, ctx, db, "gorm_test_b", "2") + + first := waitForNotification(t, listener) + second := waitForNotification(t, listener) + if first.Channel != "gorm_test_a" || first.Payload != "1" { + t.Errorf("first notification = %+v, want channel gorm_test_a payload 1", first) + } + if second.Channel != "gorm_test_b" || second.Payload != "2" { + t.Errorf("second notification = %+v, want channel gorm_test_b payload 2", second) + } +} + +// TestListener_unlisten verifies that notifications sent after Unlisten are +// not delivered. A control channel avoids timing-based assertions: PostgreSQL +// delivers notifications from different transactions in commit order, so if +// the control notification arrives first, the unlistened one was dropped. +func TestListener_unlisten(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + for _, channel := range []string{"gorm_test_dropped", "gorm_test_control"} { + if err := listener.Listen(ctx, channel); err != nil { + t.Fatalf("Listen(%q) error = %v", channel, err) + } + } + if err := listener.Unlisten(ctx, "gorm_test_dropped"); err != nil { + t.Fatalf("Unlisten() error = %v", err) + } + + notify(t, ctx, db, "gorm_test_dropped", "should not arrive") + notify(t, ctx, db, "gorm_test_control", "control") + + n := waitForNotification(t, listener) + if n.Channel != "gorm_test_control" { + t.Errorf("received notification on %q, want only %q", n.Channel, "gorm_test_control") + } +} + +func TestListener_unlistenAll(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + for _, channel := range []string{"gorm_test_one", "gorm_test_two"} { + if err := listener.Listen(ctx, channel); err != nil { + t.Fatalf("Listen(%q) error = %v", channel, err) + } + } + if err := listener.UnlistenAll(ctx); err != nil { + t.Fatalf("UnlistenAll() error = %v", err) + } + // Re-register only the control channel; notifications for the previously + // registered channels must no longer be delivered. + if err := listener.Listen(ctx, "gorm_test_control"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + + notify(t, ctx, db, "gorm_test_one", "1") + notify(t, ctx, db, "gorm_test_two", "2") + notify(t, ctx, db, "gorm_test_control", "control") + + n := waitForNotification(t, listener) + if n.Channel != "gorm_test_control" { + t.Errorf("received notification on %q, want only %q", n.Channel, "gorm_test_control") + } +} + +func TestListener_unusualChannelNames(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + + for _, channel := range []string{"my channel", "my-channel", "MyChannel", `we"ird`} { + t.Run(channel, func(t *testing.T) { + listener := newTestListener(t, ctx, db) + if err := listener.Listen(ctx, channel); err != nil { + t.Fatalf("Listen(%q) error = %v", channel, err) + } + notify(t, ctx, db, channel, "payload") + n := waitForNotification(t, listener) + if n.Channel != channel { + t.Errorf("Channel = %q, want %q", n.Channel, channel) + } + }) + } +} + +func TestListener_emptyChannel(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + if err := listener.Listen(ctx, ""); err == nil { + t.Error("Listen(\"\") succeeded, want error") + } + if err := listener.Unlisten(ctx, ""); err == nil { + t.Error("Unlisten(\"\") succeeded, want error") + } +} + +func TestListener_waitCanceled(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + if err := listener.Listen(ctx, "gorm_test_cancel"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + + waitCtx, cancel := context.WithCancel(ctx) + done := make(chan error, 1) + go func() { + _, err := listener.WaitForNotification(waitCtx) + done <- err + }() + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("WaitForNotification() error = %v, want context.Canceled", err) + } + case <-time.After(10 * time.Second): + t.Fatal("WaitForNotification() did not return after context cancellation") + } + + // A canceled wait must leave the connection and its registrations usable. + notify(t, ctx, db, "gorm_test_cancel", "after cancel") + n := waitForNotification(t, listener) + if n.Payload != "after cancel" { + t.Errorf("Payload = %q, want %q", n.Payload, "after cancel") + } + + // A deadline-exceeded wait reports context.DeadlineExceeded. + timeoutCtx, cancelTimeout := context.WithTimeout(ctx, 50*time.Millisecond) + defer cancelTimeout() + if _, err := listener.WaitForNotification(timeoutCtx); !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("WaitForNotification() error = %v, want context.DeadlineExceeded", err) + } +} + +func TestListener_customNotificationHandler(t *testing.T) { + dsn := os.Getenv("POSTGRES_DSN") + if dsn == "" { + t.Skip("POSTGRES_DSN is not set, skipping integration test") + } + + config, err := pgx.ParseConfig(dsn) + if err != nil { + t.Fatalf("pgx.ParseConfig() error = %v", err) + } + handled := make(chan *pgconn.Notification, 1) + config.OnNotification = func(_ *pgconn.PgConn, notification *pgconn.Notification) { + handled <- notification + } + + sqlDB := stdlib.OpenDB(*config) + t.Cleanup(func() { _ = sqlDB.Close() }) + db, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), &gorm.Config{}) + if err != nil { + t.Fatalf("gorm.Open() error = %v", err) + } + + ctx := context.Background() + listener := newTestListener(t, ctx, db) + if err := listener.Listen(ctx, "gorm_test_custom_handler"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + notify(t, ctx, db, "gorm_test_custom_handler", "handled") + + waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + if _, err := listener.WaitForNotification(waitCtx); !errors.Is(err, postgres.ErrNotificationHandlerConflict) { + t.Fatalf("WaitForNotification() error = %v, want ErrNotificationHandlerConflict", err) + } + + select { + case notification := <-handled: + if notification.Channel != "gorm_test_custom_handler" || notification.Payload != "handled" { + t.Errorf("handled notification = %+v, want custom-handler notification", notification) + } + default: + t.Fatal("custom OnNotification handler was not called") + } +} + +func TestListener_transactionCommitAndRollback(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + if err := listener.Listen(ctx, "gorm_test_tx"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + + // A notification sent inside a rolled-back transaction is never + // delivered. + rollbackErr := errors.New("force rollback") + err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + notify(t, ctx, tx, "gorm_test_tx", "rolled back") + return rollbackErr + }) + if !errors.Is(err, rollbackErr) { + t.Fatalf("Transaction() error = %v, want forced rollback", err) + } + + // A notification sent inside a committed transaction is delivered only + // after commit: before the commit it must not be observable. + tx := db.WithContext(ctx).Begin() + if tx.Error != nil { + t.Fatalf("Begin() error = %v", tx.Error) + } + notify(t, ctx, tx, "gorm_test_tx", "committed") + + preCommitCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + _, waitErr := listener.WaitForNotification(preCommitCtx) + cancel() + if !errors.Is(waitErr, context.DeadlineExceeded) { + t.Fatalf("WaitForNotification() before commit error = %v, want context.DeadlineExceeded", waitErr) + } + + if err := tx.Commit().Error; err != nil { + t.Fatalf("Commit() error = %v", err) + } + + n := waitForNotification(t, listener) + if n.Payload != "committed" { + t.Errorf("Payload = %q, want %q (rolled-back notification must not be delivered)", n.Payload, "committed") + } +} + +func TestListener_unrelatedChannels(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + listener := newTestListener(t, ctx, db) + + if err := listener.Listen(ctx, "gorm_test_mine"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + + notify(t, ctx, db, "gorm_test_other", "not mine") + notify(t, ctx, db, "gorm_test_mine", "mine") + + n := waitForNotification(t, listener) + if n.Channel != "gorm_test_mine" || n.Payload != "mine" { + t.Errorf("notification = %+v, want only the gorm_test_mine notification", n) + } +} + +func TestListener_closeReleasesConnection(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + + sqlDB, err := db.DB() + if err != nil { + t.Fatalf("db.DB() error = %v", err) + } + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { sqlDB.SetMaxOpenConns(0) }) + + listener, err := postgres.NewListener(ctx, db) + if err != nil { + t.Fatalf("NewListener() error = %v", err) + } + + // While the listener holds the pool's only connection, other queries + // cannot run. + busyCtx, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + busyErr := db.WithContext(busyCtx).Exec("SELECT 1").Error + cancel() + if busyErr == nil { + t.Fatal("Exec() succeeded while the listener held the pool's only connection, want timeout") + } + + if err := listener.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + + // After Close the pool slot is free again. + if err := db.WithContext(ctx).Exec("SELECT 1").Error; err != nil { + t.Fatalf("Exec() after Close error = %v", err) + } +} + +func TestListener_afterClose(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + + listener, err := postgres.NewListener(ctx, db) + if err != nil { + t.Fatalf("NewListener() error = %v", err) + } + if err := listener.Listen(ctx, "gorm_test_closed"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + if err := listener.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } + if err := listener.Close(); err != nil { + t.Fatalf("second Close() error = %v, want nil", err) + } + + if err := listener.Listen(ctx, "gorm_test_closed"); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("Listen() after Close error = %v, want ErrListenerClosed", err) + } + if err := listener.Unlisten(ctx, "gorm_test_closed"); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("Unlisten() after Close error = %v, want ErrListenerClosed", err) + } + if err := listener.Listen(ctx, ""); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("Listen(\"\") after Close error = %v, want ErrListenerClosed", err) + } + if err := listener.Unlisten(ctx, ""); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("Unlisten(\"\") after Close error = %v, want ErrListenerClosed", err) + } + if err := listener.UnlistenAll(ctx); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("UnlistenAll() after Close error = %v, want ErrListenerClosed", err) + } + if _, err := listener.WaitForNotification(ctx); !errors.Is(err, postgres.ErrListenerClosed) { + t.Errorf("WaitForNotification() after Close error = %v, want ErrListenerClosed", err) + } +} + +func TestListener_sessionAndTransactionHandles(t *testing.T) { + db := integrationDB(t) + ctx := context.Background() + + t.Run("session handle", func(t *testing.T) { + session := db.Session(&gorm.Session{}) + listener := newTestListener(t, ctx, session) + if err := listener.Listen(ctx, "gorm_test_session"); err != nil { + t.Fatalf("Listen() error = %v", err) + } + notify(t, ctx, db, "gorm_test_session", "via session") + if n := waitForNotification(t, listener); n.Payload != "via session" { + t.Errorf("Payload = %q, want %q", n.Payload, "via session") + } + }) + + // A transaction-derived handle yields an independent listener connection; + // the listener does not join the transaction. + t.Run("transaction handle", func(t *testing.T) { + err := db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + listener, err := postgres.NewListener(ctx, tx) + if err != nil { + return err + } + defer listener.Close() + if err := listener.Listen(ctx, "gorm_test_tx_handle"); err != nil { + return err + } + // Sent outside the transaction, so delivered immediately. + notify(t, ctx, db, "gorm_test_tx_handle", "independent") + n := waitForNotification(t, listener) + if n.Payload != "independent" { + t.Errorf("Payload = %q, want %q", n.Payload, "independent") + } + return nil + }) + if err != nil { + t.Fatalf("Transaction() error = %v", err) + } + }) +} diff --git a/listener_test.go b/listener_test.go new file mode 100644 index 0000000..b48becd --- /dev/null +++ b/listener_test.go @@ -0,0 +1,129 @@ +package postgres + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "testing" + + "gorm.io/gorm" +) + +func Test_quoteChannel(t *testing.T) { + tests := []struct { + name string + channel string + want string + wantErr bool + }{ + {name: "simple", channel: "events", want: `"events"`}, + {name: "uppercase", channel: "Events", want: `"Events"`}, + {name: "space", channel: "my channel", want: `"my channel"`}, + {name: "dash", channel: "my-channel", want: `"my-channel"`}, + {name: "embedded double quote", channel: `we"ird`, want: `"we""ird"`}, + {name: "asterisk is a plain identifier", channel: "*", want: `"*"`}, + {name: "empty", channel: "", wantErr: true}, + {name: "embedded NUL", channel: "my\x00channel", wantErr: true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := quoteChannel(tt.channel) + if (err != nil) != tt.wantErr { + t.Fatalf("quoteChannel() error = %v, wantErr %v", err, tt.wantErr) + } + if got != tt.want { + t.Errorf("quoteChannel() = %v, want %v", got, tt.want) + } + }) + } +} + +// fakeDriver is a database/sql driver that is not pgx, used to verify that +// NewListener rejects unsupported driver connections instead of panicking. +type fakeDriver struct{} + +func (fakeDriver) Open(name string) (driver.Conn, error) { return fakeConn{}, nil } + +type fakeConn struct{} + +func (fakeConn) Prepare(query string) (driver.Stmt, error) { + return nil, errors.New("fake driver: not implemented") +} +func (fakeConn) Close() error { return nil } +func (fakeConn) Begin() (driver.Tx, error) { + return nil, errors.New("fake driver: not implemented") +} + +func openFakeGormDB(t *testing.T) (*gorm.DB, *sql.DB) { + t.Helper() + sqlDB := sql.OpenDB(fakeConnector{}) + db, err := gorm.Open(New(Config{Conn: sqlDB}), &gorm.Config{DisableAutomaticPing: true}) + if err != nil { + t.Fatalf("gorm.Open() error = %v", err) + } + return db, sqlDB +} + +type fakeConnector struct{} + +func (fakeConnector) Connect(context.Context) (driver.Conn, error) { return fakeConn{}, nil } +func (fakeConnector) Driver() driver.Driver { return fakeDriver{} } + +func TestNewListener_nilDB(t *testing.T) { + if _, err := NewListener(context.Background(), nil); !errors.Is(err, gorm.ErrInvalidDB) { + t.Fatalf("NewListener(nil) error = %v, want gorm.ErrInvalidDB", err) + } +} + +func TestNewListener_unsupportedDriver(t *testing.T) { + db, _ := openFakeGormDB(t) + if _, err := NewListener(context.Background(), db); !errors.Is(err, ErrUnsupportedConnection) { + t.Fatalf("NewListener() error = %v, want ErrUnsupportedConnection", err) + } +} + +func TestNewListener_invalidConnPool(t *testing.T) { + db, _ := openFakeGormDB(t) + // A ConnPool that is neither *sql.DB nor a GetDBConnector cannot provide + // a listener connection. + db.ConnPool = fakeConnPool{} + db.Statement.ConnPool = fakeConnPool{} + if _, err := NewListener(context.Background(), db); !errors.Is(err, gorm.ErrInvalidDB) { + t.Fatalf("NewListener() error = %v, want gorm.ErrInvalidDB", err) + } +} + +type fakeConnPool struct{} + +func (fakeConnPool) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) { + return nil, errors.New("fake pool: not implemented") +} +func (fakeConnPool) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { + return nil, errors.New("fake pool: not implemented") +} +func (fakeConnPool) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { + return nil, errors.New("fake pool: not implemented") +} +func (fakeConnPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { + return nil +} + +func TestNewListener_closedDB(t *testing.T) { + db, sqlDB := openFakeGormDB(t) + if err := sqlDB.Close(); err != nil { + t.Fatalf("sqlDB.Close() error = %v", err) + } + if _, err := NewListener(context.Background(), db); err == nil { + t.Fatal("NewListener() on closed database succeeded, want error") + } +} + +func TestNewListener_canceledContext(t *testing.T) { + db, _ := openFakeGormDB(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := NewListener(ctx, db); !errors.Is(err, context.Canceled) { + t.Fatalf("NewListener() error = %v, want context.Canceled", err) + } +}