Set the async_status to ASYNC_DONE after a query with wait callback.

Failing in doing that broke notifications reception.

The responsibility for changing the async_status has been moved to the
poll function: this is consistent with how the async branch is
implemented.

With this commit all the test suite passes in "green" mode.

2 files changed, 13 insertions, 14 deletions

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index e564cfc..e12e3c3 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -721,6 +721,9 @@ conn_poll_green(connectionObject *self)
             if (PQisBusy(self->pgconn)) {
                 res = PSYCO_POLL_READ;
             } else {
+                /* Reading complete: set the async status so that a spare poll()
+                   will only look for NOTIFYs */
+                self->async_status = ASYNC_DONE;
                 res = PSYCO_POLL_OK;
             }
             break;
@@ -729,7 +732,10 @@ conn_poll_green(connectionObject *self)
             Dprintf("conn_poll: async_status = ASYNC_WRITE");
             switch (PQflush(self->pgconn)) {
             case  0:  /* success */
-                res = PSYCO_POLL_OK;
+                /* we've finished pushing the query to the server. Let's start
+                  reading the results. */
+                self->async_status = ASYNC_READ;
+                res = PSYCO_POLL_READ;
                 break;
             case  1:  /* would block */
                 res = PSYCO_POLL_WRITE;
diff --git a/psycopg/green.c b/psycopg/green.c
index df1c02e..9e7e5d2 100644
--- a/psycopg/green.c
+++ b/psycopg/green.c
@@ -158,23 +158,15 @@ psyco_exec_green(connectionObject *conn, const char *command)
         goto clear;
     }
 
-    /* Ensure the query reached the server. */
+    /* Enter the poll loop with a write. When writing is finished the poll
+       implementation will set the status to ASYNC_READ without exiting the
+       loop. If read is finished the status is finally set to ASYNC_DONE.
+    */
     conn->async_status = ASYNC_WRITE;
 
     pyrv = PyObject_CallFunctionObjArgs(cb, conn, NULL);
     if (!pyrv) {
-        Dprintf("psyco_exec_green: error in callback sending query");
-        psyco_clear_result_blocking(conn);
-        goto clear;
-    }
-    Py_DECREF(pyrv);
-
-    /* Loop reading data using the user-provided wait function */
-    conn->async_status = ASYNC_READ;
-
-    pyrv = PyObject_CallFunctionObjArgs(cb, conn, NULL);
-    if (!pyrv) {
-        Dprintf("psyco_exec_green: error in callback reading result");
+        Dprintf("psyco_exec_green: error in wait callback");
         psyco_clear_result_blocking(conn);
         goto clear;
     }
@@ -195,6 +187,7 @@ psyco_exec_green(connectionObject *conn, const char *command)
     }
 
 clear:
+    conn->async_status = ASYNC_DONE;
     Py_DECREF(cb);
 end:
     return result;