Add curs_get_last_result, a function to get the last result from a connection

3 files changed, 38 insertions, 13 deletions

diff --git a/psycopg/cursor.h b/psycopg/cursor.h
index 0fd1c69..a422283 100644
--- a/psycopg/cursor.h
+++ b/psycopg/cursor.h
@@ -86,6 +86,7 @@ typedef struct {
 
 /* C-callable functions in cursor_int.c and cursor_ext.c */
 HIDDEN void curs_reset(cursorObject *self);
+HIDDEN void curs_get_last_result(cursorObject *self);
 
 /* exception-raising macros */
 #define EXC_IF_CURS_CLOSED(self) \
diff --git a/psycopg/cursor_int.c b/psycopg/cursor_int.c
index 004ba82..ae28af1 100644
--- a/psycopg/cursor_int.c
+++ b/psycopg/cursor_int.c
@@ -54,3 +54,28 @@ curs_reset(cursorObject *self)
     self->casts = NULL;
     Py_XDECREF(tmp);
 }
+
+/*
+ * curs_get_last_result
+ *
+ * read all results from the connection, save the last one
+ */
+
+void
+curs_get_last_result(cursorObject *self) {
+    PGresult *pgres;
+
+    IFCLEARPGRES(self->pgres);
+    Py_BEGIN_ALLOW_THREADS;
+    pthread_mutex_lock(&(self->conn->lock));
+    /* read all results: there can be multiple if the client sent multiple
+       statements */
+    while ((pgres = PQgetResult(self->conn->pgconn)) != NULL) {
+        IFCLEARPGRES(self->pgres);
+        self->pgres = pgres;
+    }
+    self->conn->async_cursor = NULL;
+    pthread_mutex_unlock(&(self->conn->lock));
+    Py_END_ALLOW_THREADS;
+    self->needsfetch = 1;
+}
diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c
index 4c66a06..3bcf0e4 100644
--- a/psycopg/cursor_type.c
+++ b/psycopg/cursor_type.c
@@ -656,20 +656,19 @@ _psyco_curs_prefetch(cursorObject *self)
 {
     int i = 0;
 
-    /* check if the fetching cursor is the one that did the asynchronous query
-       and raise an exception if not */
-    Py_BEGIN_ALLOW_THREADS;
-    pthread_mutex_lock(&(self->conn->lock));
-    if (self->conn->async_cursor != NULL
-        && self->conn->async_cursor != (PyObject*)self) {
-        pthread_mutex_unlock(&(self->conn->lock));
-        Py_BLOCK_THREADS;
-        psyco_set_error(ProgrammingError, (PyObject*)self,
-                         "asynchronous fetch by wrong cursor", NULL, NULL);
-        return -2;
+    /* check if there is an asynchronous query in progess and block until data
+       is read from it */
+    if (self->conn->async_cursor) {
+        /* first check if it's the right cursor */
+        if (self->conn->async_cursor != (PyObject*)self) {
+            psyco_set_error(ProgrammingError, (PyObject*)self,
+                            "asynchronous fetch by wrong cursor", NULL, NULL);
+            return -2;
+        }
+        /* now get the result */
+        Dprintf("_psyco_curs_prefetch: blocking until all data is read");
+        curs_get_last_result(self);
     }
-    pthread_mutex_unlock(&(self->conn->lock));
-    Py_END_ALLOW_THREADS;
 
     if (self->pgres == NULL || self->needsfetch) {
         self->needsfetch = 0;