Implement dblink_get_notify().

Adds the ability to retrieve async notifications using dblink,
via the addition of the function dblink_get_notify(). Original patch
by Marcus Kempe, suggestions by Tom Lane and Alvaro Herrera, patch
review and adjustments by Joe Conway.

7 files changed, 248 insertions, 0 deletions

diff --git a/contrib/dblink/dblink.c b/contrib/dblink/dblink.c
index 276c7e1400..83b334c4f6 100644
--- a/contrib/dblink/dblink.c
+++ b/contrib/dblink/dblink.c
@@ -1635,6 +1635,89 @@ dblink_current_query(PG_FUNCTION_ARGS)
 	PG_RETURN_DATUM(current_query(fcinfo));
 }
 
+/*
+ * Retrieve async notifications for a connection. 
+ *
+ * Returns an setof record of notifications, or an empty set if none recieved.
+ * Can optionally take a named connection as parameter, but uses the unnamed connection per default.
+ *
+ */
+#define DBLINK_NOTIFY_COLS		3
+
+PG_FUNCTION_INFO_V1(dblink_get_notify);
+Datum
+dblink_get_notify(PG_FUNCTION_ARGS)
+{
+	PGconn			   *conn = NULL;
+	remoteConn		   *rconn = NULL;
+	PGnotify		   *notify;
+	ReturnSetInfo	   *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	TupleDesc			tupdesc;
+	Tuplestorestate	   *tupstore;
+	MemoryContext		per_query_ctx;
+	MemoryContext		oldcontext;
+
+	DBLINK_INIT;
+	if (PG_NARGS() == 1)
+		DBLINK_GET_NAMED_CONN;
+	else
+		conn = pconn->conn;
+
+	/* create the tuplestore */
+	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
+	oldcontext = MemoryContextSwitchTo(per_query_ctx);
+
+	tupdesc = CreateTemplateTupleDesc(DBLINK_NOTIFY_COLS, false);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 1, "notify_name",
+					   TEXTOID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 2, "be_pid",
+					   INT4OID, -1, 0);
+	TupleDescInitEntry(tupdesc, (AttrNumber) 3, "extra",
+					   TEXTOID, -1, 0);
+
+	tupstore = tuplestore_begin_heap(true, false, work_mem);
+	rsinfo->returnMode = SFRM_Materialize;
+	rsinfo->setResult = tupstore;
+	rsinfo->setDesc = tupdesc;
+
+	MemoryContextSwitchTo(oldcontext);
+
+	PQconsumeInput(conn);
+	while ((notify = PQnotifies(conn)) != NULL)
+	{
+		Datum		values[DBLINK_NOTIFY_COLS];
+		bool		nulls[DBLINK_NOTIFY_COLS];
+
+		memset(values, 0, sizeof(values));
+		memset(nulls, 0, sizeof(nulls));
+
+		if (notify->relname != NULL)
+			values[0] = CStringGetTextDatum(notify->relname);
+		else
+			nulls[0] = true;
+
+		values[1] = Int32GetDatum(notify->be_pid);
+
+		if (notify->extra != NULL)
+			values[2] = CStringGetTextDatum(notify->extra);
+		else
+			nulls[2] = true;
+
+		/* switch to appropriate context while storing the tuple */
+		MemoryContextSwitchTo(per_query_ctx);
+		tuplestore_putvalues(tupstore, tupdesc, values, nulls);
+		MemoryContextSwitchTo(oldcontext);
+
+		PQfreemem(notify);
+		PQconsumeInput(conn);
+	}
+
+	/* clean up and return the tuplestore */
+	tuplestore_donestoring(tupstore);
+
+	return (Datum) 0;
+}
+
 /*************************************************************
  * internal functions
  */
diff --git a/contrib/dblink/dblink.h b/contrib/dblink/dblink.h
index 7dab7eed86..255f5d011a 100644
--- a/contrib/dblink/dblink.h
+++ b/contrib/dblink/dblink.h
@@ -57,5 +57,6 @@ extern Datum dblink_build_sql_insert(PG_FUNCTION_ARGS);
 extern Datum dblink_build_sql_delete(PG_FUNCTION_ARGS);
 extern Datum dblink_build_sql_update(PG_FUNCTION_ARGS);
 extern Datum dblink_current_query(PG_FUNCTION_ARGS);
+extern Datum dblink_get_notify(PG_FUNCTION_ARGS);
 
 #endif   /* DBLINK_H */
diff --git a/contrib/dblink/dblink.sql.in b/contrib/dblink/dblink.sql.in
index c72a87e458..da5dd65420 100644
--- a/contrib/dblink/dblink.sql.in
+++ b/contrib/dblink/dblink.sql.in
@@ -202,3 +202,22 @@ CREATE OR REPLACE FUNCTION dblink_error_message(text)
 RETURNS text
 AS 'MODULE_PATHNAME', 'dblink_error_message'
 LANGUAGE C STRICT;
+
+CREATE OR REPLACE FUNCTION dblink_get_notify(
+    OUT notify_name TEXT,
+    OUT be_pid INT4,
+    OUT extra TEXT
+) 
+RETURNS setof record
+AS 'MODULE_PATHNAME', 'dblink_get_notify'
+LANGUAGE C STRICT;
+
+CREATE OR REPLACE FUNCTION dblink_get_notify(
+    conname TEXT,
+    OUT notify_name TEXT,
+    OUT be_pid INT4,
+    OUT extra TEXT
+) 
+RETURNS setof record
+AS 'MODULE_PATHNAME', 'dblink_get_notify'
+LANGUAGE C STRICT;
diff --git a/contrib/dblink/expected/dblink.out b/contrib/dblink/expected/dblink.out
index 102444d88c..13e4a6fb4c 100644
--- a/contrib/dblink/expected/dblink.out
+++ b/contrib/dblink/expected/dblink.out
@@ -827,3 +827,54 @@ DROP USER dblink_regression_test;
 DROP USER MAPPING FOR public SERVER fdtest;
 DROP SERVER fdtest;
 DROP FOREIGN DATA WRAPPER postgresql;
+-- test asynchronous notifications
+SELECT dblink_connect('dbname=contrib_regression');
+ dblink_connect
+----------------
+ OK
+(1 row)
+
+--should return listen
+SELECT dblink_exec('LISTEN regression');
+ dblink_exec
+-------------
+ LISTEN
+(1 row)
+
+--should return listen
+SELECT dblink_exec('LISTEN foobar');
+ dblink_exec
+-------------
+ LISTEN
+(1 row)
+
+SELECT dblink_exec('NOTIFY regression');
+ dblink_exec
+-------------
+ NOTIFY
+(1 row)
+
+SELECT dblink_exec('NOTIFY foobar');
+ dblink_exec
+-------------
+ NOTIFY
+(1 row)
+
+SELECT notify_name, be_pid = (select t.be_pid from dblink('select pg_backend_pid()') as t(be_pid int)) AS is_self_notify, extra from dblink_get_notify();
+ notify_name | is_self_notify | extra
+-------------+----------------+-------
+ regression  | t              |
+ foobar      | t              |
+(2 rows)
+
+SELECT * from dblink_get_notify();
+ notify_name | be_pid | extra
+-------------+--------+-------
+(0 rows)
+
+SELECT dblink_disconnect();
+ dblink_disconnect
+-------------------
+ OK
+(1 row)
+
diff --git a/contrib/dblink/sql/dblink.sql b/contrib/dblink/sql/dblink.sql
index a8190c6564..d0ad87695a 100644
--- a/contrib/dblink/sql/dblink.sql
+++ b/contrib/dblink/sql/dblink.sql
@@ -389,3 +389,20 @@ DROP USER dblink_regression_test;
 DROP USER MAPPING FOR public SERVER fdtest;
 DROP SERVER fdtest;
 DROP FOREIGN DATA WRAPPER postgresql;
+
+-- test asynchronous notifications
+SELECT dblink_connect('dbname=contrib_regression');
+
+--should return listen
+SELECT dblink_exec('LISTEN regression');
+--should return listen
+SELECT dblink_exec('LISTEN foobar');
+
+SELECT dblink_exec('NOTIFY regression');
+SELECT dblink_exec('NOTIFY foobar');
+
+SELECT notify_name, be_pid = (select t.be_pid from dblink('select pg_backend_pid()') as t(be_pid int)) AS is_self_notify, extra from dblink_get_notify();
+
+SELECT * from dblink_get_notify();
+
+SELECT dblink_disconnect();
diff --git a/contrib/dblink/uninstall_dblink.sql b/contrib/dblink/uninstall_dblink.sql
index fa16e3dbd6..45cf13c9d0 100644
--- a/contrib/dblink/uninstall_dblink.sql
+++ b/contrib/dblink/uninstall_dblink.sql
@@ -76,3 +76,7 @@ DROP FUNCTION dblink_get_result(text, boolean);
 DROP FUNCTION dblink_is_busy(text);
 
 DROP FUNCTION dblink_send_query(text, text);
+
+DROP FUNCTION dblink_get_notify();
+
+DROP FUNCTION dblink_get_notify(text);
diff --git a/doc/src/sgml/dblink.sgml b/doc/src/sgml/dblink.sgml
index dd908a2388..0ce2168311 100644
--- a/doc/src/sgml/dblink.sgml
+++ b/doc/src/sgml/dblink.sgml
@@ -1260,6 +1260,79 @@ SELECT *
   </refsect1>
  </refentry>
 
+ <refentry id="CONTRIB-DBLINK-GET-NOTIFY">
+  <refnamediv>
+   <refname>dblink_get_notify</refname>
+   <refpurpose>retrieve async notifications on a connection</refpurpose>
+  </refnamediv>
+
+  <refsynopsisdiv>
+   <synopsis>
+    dblink_get_notify() returns setof (notify_name text, be_pid int, extra text)
+    dblink_get_notify(text connname) returns setof (notify_name text, be_pid int, extra text)
+   </synopsis>
+  </refsynopsisdiv>
+
+  <refsect1>
+   <title>Description</title>
+
+   <para>
+    <function>dblink_get_notify</> retrieves notifications on either 
+    the unnamed connection, or on a named connection if specified.
+    To receive notifications via dblink, <function>LISTEN</> must 
+    first be issued, using <function>dblink_exec</>.
+    For details see <xref linkend="sql-listen"> and <xref linkend="sql-notify">.
+   </para>
+
+  </refsect1>
+
+  <refsect1>
+   <title>Arguments</title>
+
+   <variablelist>
+    <varlistentry>
+     <term><parameter>conname</parameter></term>
+     <listitem>
+      <para>
+       The name of a named connection to get notifications on.
+      </para>
+     </listitem>
+    </varlistentry>
+   </variablelist>
+  </refsect1>
+
+  <refsect1>
+   <title>Return Value</title>
+    <para>Returns setof (notify_name text, be_pid int, extra text), or an empty set if none.</para>
+  </refsect1>
+
+  <refsect1>
+   <title>Example</title>
+
+   <programlisting>
+test=# SELECT dblink_exec('LISTEN virtual');
+ dblink_exec 
+-------------
+ LISTEN
+(1 row)
+
+test=# SELECT * FROM dblink_get_notify();
+ notify_name | be_pid | extra
+-------------+--------+-------
+(0 rows)
+
+test=# NOTIFY virtual;
+NOTIFY
+
+SELECT * FROM dblink_get_notify();
+ notify_name | be_pid | extra
+-------------+--------+-------
+ virtual     |   1229 |
+(1 row)
+   </programlisting>
+  </refsect1>
+ </refentry>
+
  <refentry id="CONTRIB-DBLINK-GET-RESULT">
   <refmeta>
    <refentrytitle>dblink_get_result</refentrytitle>