db.pg: support LISTEN/NOTIFY (implement feature #23983) (#26133)

fleximus · web-flow · commit 3f75ccf50560 · 2025-12-26T09:03:17.000+02:00
diff --git a/vlib/db/pg/README.md b/vlib/db/pg/README.md
@@ -100,3 +100,81 @@ The number following the $ specifies which parameter from the argument array to
 db.exec_param_many('INSERT INTO users (username, password) VALUES ($1, $2)', ['tom', 'securePassword']) or { panic(err) }
 db.exec_param('SELECT * FROM users WHERE username = ($1) limit 1', 'tom') or { panic(err) }
 ```
+
+## Using LISTEN/NOTIFY
+
+PostgreSQL's LISTEN/NOTIFY mechanism allows you to build event-driven applications. One
+connection can send notifications on a channel, and all connections listening on that
+channel will receive them.
+
+### Basic Usage
+
+```v ignore
+import db.pg
+
+fn main() {
+	db := pg.connect(pg.Config{ user: 'postgres', password: 'password', dbname: 'mydb' })!
+	defer { db.close() or {} }
+
+	// Start listening on a channel
+	db.listen('my_channel')!
+
+	// From another connection or session, send a notification
+	db.notify('my_channel', 'Hello, World!')!
+
+	// Process incoming data from the server
+	db.consume_input()!
+
+	// Check for notifications
+	if notification := db.get_notification() {
+		println('Received notification on channel: ${notification.channel}')
+		println('Payload: ${notification.payload}')
+		println('From server process: ${notification.pid}')
+	}
+
+	// Stop listening
+	db.unlisten('my_channel')!
+	// Or unlisten from all channels
+	db.unlisten_all()!
+}
+```
+
+### Event Loop with Polling
+
+For real-time applications, you can use the socket file descriptor with select/poll:
+
+```v ignore
+import db.pg
+import time
+
+fn main() {
+	db := pg.connect(pg.Config{ user: 'postgres', password: 'password', dbname: 'mydb' })!
+	defer { db.close() or {} }
+
+	db.listen('events')!
+
+	// Get socket fd for polling (useful with select/epoll)
+	socket_fd := db.socket()
+	println('Socket FD: ${socket_fd}')
+
+	// Simple polling loop
+	for {
+		db.consume_input()!
+		for {
+			notification := db.get_notification() or { break }
+			println('Event: ${notification.channel} - ${notification.payload}')
+		}
+		time.sleep(100 * time.millisecond)
+	}
+}
+```
+
+### Available Methods
+
+- `listen(channel string)` - Register to receive notifications on a channel
+- `unlisten(channel string)` - Unregister from a specific channel
+- `unlisten_all()` - Unregister from all channels
+- `notify(channel string, payload string)` - Send a notification (payload can be empty)
+- `consume_input()` - Read pending data from server (call before get_notification)
+- `get_notification()` - Returns the next pending notification, or none if there are no notifications
+- `socket()` - Returns the connection's socket file descriptor for use with select/poll
diff --git a/vlib/db/pg/pg.c.v b/vlib/db/pg/pg.c.v
@@ -72,6 +72,14 @@ pub:
 	rows []Row
 }
 
+// Notification represents a notification received from the server via LISTEN/NOTIFY
+pub struct Notification {
+pub:
+	channel string // notification channel name
+	pid     int    // process ID of notifying server process
+	payload string // notification payload string (may be empty)
+}
+
 pub struct Config {
 pub:
 	host     string = 'localhost'
@@ -93,6 +101,14 @@ pub struct C.PGresult {}
 @[typedef]
 pub struct C.PGconn {}
 
+// PGnotify represents a notification received from the server via LISTEN/NOTIFY
+@[typedef]
+pub struct C.PGnotify {
+	relname &char // notification channel name
+	be_pid  int   // process ID of notifying server process
+	extra   &char // notification payload string
+}
+
 pub enum ConnStatusType {
 	ok  = C.CONNECTION_OK
 	bad = C.CONNECTION_BAD
@@ -177,6 +193,15 @@ fn C.PQfreemem(ptr voidptr)
 
 fn C.PQfinish(conn &C.PGconn)
 
+// LISTEN/NOTIFY support
+fn C.PQnotifies(conn &C.PGconn) &C.PGnotify
+
+fn C.PQconsumeInput(conn &C.PGconn) int
+
+fn C.PQsocket(conn &C.PGconn) int
+
+fn C.PQescapeLiteral(conn &C.PGconn, str &char, length usize) &char
+
 // connect makes a new connection to the database server using
 // the parameters from the `Config` structure, returning
 // a connection error when something goes wrong
@@ -614,3 +639,103 @@ pub fn (res Result) as_structs[T](mapper fn (Result, Row) !T) ![]T {
 
 	return typed
 }
+
+// listen registers the connection to receive notifications on the specified channel.
+// After calling this, use consume_input() and get_notification() to receive notifications.
+pub fn (db &DB) listen(channel string) ! {
+	if !channel.is_identifier() {
+		return error('channel name should be a valid identifier')
+	}
+	sql_stmt := 'LISTEN ${channel};'
+	_ := C.PQexec(db.conn, &char(sql_stmt.str))
+	e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+	if e != '' {
+		return error('pg listen error: "${e}"')
+	}
+}
+
+// unlisten unregisters the connection from receiving notifications on the specified channel.
+// Use unlisten_all() to unregister from all channels.
+pub fn (db &DB) unlisten(channel string) ! {
+	if !channel.is_identifier() {
+		return error('channel name should be a valid identifier')
+	}
+	sql_stmt := 'UNLISTEN ${channel};'
+	_ := C.PQexec(db.conn, &char(sql_stmt.str))
+	e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+	if e != '' {
+		return error('pg unlisten error: "${e}"')
+	}
+}
+
+// unlisten_all unregisters the connection from all notification channels.
+pub fn (db &DB) unlisten_all() ! {
+	_ := C.PQexec(db.conn, c'UNLISTEN *;')
+	e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+	if e != '' {
+		return error('pg unlisten error: "${e}"')
+	}
+}
+
+// notify sends a notification on the specified channel with an optional payload.
+// All connections currently listening on that channel will receive the notification.
+pub fn (db &DB) notify(channel string, payload string) ! {
+	if !channel.is_identifier() {
+		return error('channel name should be a valid identifier')
+	}
+	mut sql_stmt := ''
+	if payload.len > 0 {
+		// Use PQescapeLiteral to safely escape the payload
+		escaped := C.PQescapeLiteral(db.conn, &char(payload.str), usize(payload.len))
+		if escaped == unsafe { nil } {
+			e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+			return error('pg notify error: failed to escape payload: "${e}"')
+		}
+		sql_stmt = unsafe { 'NOTIFY ${channel}, ' + escaped.vstring() + ';' }
+		C.PQfreemem(escaped)
+	} else {
+		sql_stmt = 'NOTIFY ${channel};'
+	}
+	_ := C.PQexec(db.conn, &char(sql_stmt.str))
+	e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+	if e != '' {
+		return error('pg notify error: "${e}"')
+	}
+}
+
+// consume_input reads any available input from the server.
+// This must be called before get_notification() to ensure pending notifications are processed.
+// Returns true on success, false if there was an error reading from the connection.
+pub fn (db &DB) consume_input() !bool {
+	result := C.PQconsumeInput(db.conn)
+	if result == 0 {
+		e := unsafe { C.PQerrorMessage(db.conn).vstring() }
+		return error('pg consume_input error: "${e}"')
+	}
+	return true
+}
+
+// get_notification returns the next pending notification from the server, if any.
+// Returns none if there are no pending notifications.
+// You should call consume_input() before this to ensure all pending notifications are available.
+pub fn (db &DB) get_notification() ?Notification {
+	notify := C.PQnotifies(db.conn)
+	if notify == unsafe { nil } {
+		return none
+	}
+	defer {
+		C.PQfreemem(notify)
+	}
+	return Notification{
+		channel: unsafe { notify.relname.vstring() }
+		pid:     notify.be_pid
+		payload: unsafe { notify.extra.vstring() }
+	}
+}
+
+// socket returns the file descriptor of the connection socket to the server.
+// This is useful for applications that want to use select() or poll() to wait
+// for notifications without blocking. Returns -1 if no valid socket.
+pub fn (db &DB) socket() int {
+	return C.PQsocket(db.conn)
+}
diff --git a/vlib/db/pg/pg_test.v b/vlib/db/pg/pg_test.v
@@ -80,3 +80,60 @@ fn test_transaction() {
 		dump(row.str())
 	}
 }
+
+fn test_listen_notify() {
+	$if !network ? {
+		eprintln('> Skipping test ${@FN}, since `-d network` is not passed.')
+		eprintln('> This test requires a working postgres server running on localhost.')
+		return
+	}
+
+	db := pg.connect(pg.Config{ user: 'postgres', password: '12345678', dbname: 'postgres' })!
+	defer {
+		db.close() or {}
+	}
+
+	// Test listen
+	db.listen('test_channel')!
+
+	// Test notify with payload
+	db.notify('test_channel', 'hello world')!
+
+	// Consume input to process the notification
+	db.consume_input()!
+
+	// Get the notification
+	if notification := db.get_notification() {
+		assert notification.channel == 'test_channel'
+		assert notification.payload == 'hello world'
+		assert notification.pid > 0
+	} else {
+		assert false, 'Expected a notification but got none'
+	}
+
+	// Test notify without payload
+	db.notify('test_channel', '')!
+	db.consume_input()!
+
+	if notification := db.get_notification() {
+		assert notification.channel == 'test_channel'
+		assert notification.payload == ''
+	} else {
+		assert false, 'Expected a notification but got none'
+	}
+
+	// Test that no more notifications are pending
+	assert db.get_notification() == none
+
+	// Test unlisten
+	db.unlisten('test_channel')!
+
+	// Test unlisten_all
+	db.listen('channel1')!
+	db.listen('channel2')!
+	db.unlisten_all()!
+
+	// Test socket (should return valid fd)
+	socket_fd := db.socket()
+	assert socket_fd >= 0
+}
