Support query cancellation.

Add a cancel() method do the connection object that will interrupt
the current query using the libpq PQcancel() function.

1 files changed, 19 insertions, 0 deletions

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index f98e4fd..73292b8 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -282,6 +282,12 @@ conn_get_server_version(PGconn *pgconn)
     return (int)PQserverVersion(pgconn);
 }
 
+PGcancel *
+conn_get_cancel(PGconn *pgconn)
+{
+    return PQgetCancel(pgconn);
+}
+
 
 /* Return 1 if the server datestyle allows us to work without problems,
    0 if it needs to be set to something better, e.g. ISO. */
@@ -320,6 +326,12 @@ conn_setup(connectionObject *self, PGconn *pgconn)
         return -1;
     }
 
+    self->cancel = conn_get_cancel(self->pgconn);
+    if (self->cancel == NULL) {
+        PyErr_SetString(OperationalError, "can't get cancellation key");
+        return -1;
+    }
+
     Py_BEGIN_ALLOW_THREADS;
     pthread_mutex_lock(&self->lock);
     Py_BLOCK_THREADS;
@@ -645,6 +657,11 @@ _conn_poll_setup_async(connectionObject *self)
         if (self->encoding == NULL) {
             break;
         }
+        self->cancel = conn_get_cancel(self->pgconn);
+        if (self->cancel == NULL) {
+            PyErr_SetString(OperationalError, "can't get cancellation key");
+            break;
+        }
 
         /* asynchronous connections always use isolation level 0, the user is
          * expected to manage the transactions himself, by sending
@@ -782,8 +799,10 @@ conn_close(connectionObject *self)
 
     if (self->pgconn) {
         PQfinish(self->pgconn);
+        PQfreeCancel(self->cancel);
         Dprintf("conn_close: PQfinish called");
         self->pgconn = NULL;
+        self->cancel = NULL;
     }
 
     pthread_mutex_unlock(&self->lock);