If query is canceled by user, send cancel request

to remote databases too.

Patch by Ye Wenbin

diff --git a/src/execute.c b/src/execute.c
index a8b8cd05d17662a48c91de4c87d4ce2d8ab22241..385677ac08fe432ca13855455340671468d72a8e 100644 (file)
--- a/src/execute.c
+++ b/src/execute.c
@@ -31,6 +31,16 @@
 
 #include "poll_compat.h"
 
+#if PG_VERSION_NUM < 80400
+static int geterrcode(void)
+{
+       ErrorData *edata = CopyErrorData();
+       int code = edata->sqlerrcode;
+       FreeErrorData(edata); /* unnecessary? */
+       return code;
+}
+#endif
+
 /* some error happened */
 static void
 conn_error(ProxyFunction *func, ProxyConnection *conn, const char *desc)
@@ -613,6 +623,40 @@ remote_execute(ProxyFunction *func,
        }
 }
 
+static void
+remote_cancel(ProxyFunction *func)
+{
+       ProxyConnection *conn;
+       ProxyCluster *cluster = func->cur_cluster;
+       PGcancel *cancel;
+       char errbuf[256];
+       int ret;
+       int i;
+
+       if (cluster == NULL)
+               return;
+
+       for (i = 0; i < cluster->conn_count; i++)
+       {
+               conn = &cluster->conn_list[i];
+               if (conn->state == C_NONE ||
+                       conn->state == C_READY ||
+                       conn->state == C_DONE)
+                       continue;
+
+               cancel = PQgetCancel(conn->db);
+               if (cancel == NULL)
+               {
+                       elog(NOTICE, "Invalid connection!");
+                       continue;
+               }
+               ret = PQcancel(cancel, errbuf, sizeof(errbuf));
+               PQfreeCancel(cancel);
+               if (ret == 0)
+                       elog(NOTICE, "Cancel query failed!");
+       }
+}
+
 /* Run hash function and tag connections */
 static void
 tag_hash_partitions(ProxyFunction *func, FunctionCallInfo fcinfo)
@@ -751,8 +795,21 @@ plproxy_exec(ProxyFunction *func, FunctionCallInfo fcinfo)
                }
        }
 
-       if (gotbin)
-               remote_execute(func, values, plengths, pformats);
-       else
-               remote_execute(func, values, NULL, NULL);
+       /*
+        * Run query.  On cancel, send cancel request to partitions too.
+        */
+       PG_TRY();
+       {
+               if (gotbin)
+                       remote_execute(func, values, plengths, pformats);
+               else
+                       remote_execute(func, values, NULL, NULL);
+       }
+       PG_CATCH();
+       {
+               if (geterrcode() == ERRCODE_QUERY_CANCELED)
+                       remote_cancel(func);
+               PG_RE_THROW();
+       }
+       PG_END_TRY();
 }