From 7af0bf0b54524378b6fc1f58781e65b252ed730b Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 21 Apr 2010 16:20:16 +0100
Subject: [PATCH] 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.
---
 psycopg/connection_int.c |  8 +++++++-
 psycopg/green.c          | 19 ++++++-------------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index e564cfc7..e12e3c36 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 df1c02e4..9e7e5d2f 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;