From 2596cf712554833e480011de0c1b3074656ec8f3 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sun, 25 Apr 2010 23:01:53 +0100
Subject: [PATCH] Simplified psyco_wait() function interface.

---
 psycopg/connection_int.c |  6 +-----
 psycopg/green.c          | 33 ++++++++++++++-------------------
 psycopg/green.h          |  2 +-
 3 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index b8787491..3684cf47 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -346,7 +346,6 @@ static int
 _conn_sync_connect(connectionObject *self)
 {
     PGconn *pgconn;
-    PyObject *wait_rv;
     int green;
 
     /* store this value to prevent inconsistencies due to a change
@@ -382,10 +381,7 @@ _conn_sync_connect(connectionObject *self)
 
     /* if the connection is green, wait to finish connection */
     if (green) {
-        wait_rv = psyco_wait(self);
-        if (wait_rv) {
-            Py_DECREF(wait_rv);
-        } else {
+        if (0 != psyco_wait(self)) {
             return -1;
         }
     }
diff --git a/psycopg/green.c b/psycopg/green.c
index bac281a0..0036d2e5 100644
--- a/psycopg/green.c
+++ b/psycopg/green.c
@@ -114,9 +114,9 @@ have_wait_callback()
  * raise `InterfaceError` if it is not. Use `psyco_green()` to check if
  * the function is to be called.
  *
- * The function returns the return value of the called function.
+ * Return 0 on success, else nonzero and set a Python exception.
  */
-PyObject *
+int
 psyco_wait(connectionObject *conn)
 {
     PyObject *rv;
@@ -124,13 +124,19 @@ psyco_wait(connectionObject *conn)
 
     Dprintf("psyco_wait");
     if (!(cb = have_wait_callback())) {
-        return NULL;
+        return -1;
     }
 
     rv = PyObject_CallFunctionObjArgs(cb, conn, NULL);
     Py_DECREF(cb);
 
-    return rv;
+    if (NULL != rv) {
+        Py_DECREF(rv);
+        return 0;
+    } else {
+        Dprintf("psyco_wait: error in wait callback");
+        return -1;
+    }
 }
 
 /* Replacement for PQexec using the user-provided wait function.
@@ -145,16 +151,10 @@ PGresult *
 psyco_exec_green(connectionObject *conn, const char *command)
 {
     PGresult *result = NULL;
-    PyObject *cb, *pyrv;
-
-    Dprintf("psyco_exec_green: executing query async");
-    if (!(cb = have_wait_callback())) {
-        goto end;
-    }
 
     /* Send the query asynchronously */
     if (0 == pq_send_query(conn, command)) {
-        goto clear;
+        goto end;
     }
 
     /* Enter the poll loop with a write. When writing is finished the poll
@@ -163,21 +163,16 @@ psyco_exec_green(connectionObject *conn, const char *command)
     */
     conn->async_status = ASYNC_WRITE;
 
-    pyrv = PyObject_CallFunctionObjArgs(cb, conn, NULL);
-    if (!pyrv) {
-        Dprintf("psyco_exec_green: error in wait callback");
+    if (0 != psyco_wait(conn)) {
         psyco_clear_result_blocking(conn);
-        goto clear;
+        goto end;
     }
-    Py_DECREF(pyrv);
 
     /* Now we can read the data without fear of blocking. */
     result = pq_get_last_result(conn);
 
-clear:
-    conn->async_status = ASYNC_DONE;
-    Py_DECREF(cb);
 end:
+    conn->async_status = ASYNC_DONE;
     return result;
 }
 
diff --git a/psycopg/green.h b/psycopg/green.h
index 175f0552..6466ae7c 100644
--- a/psycopg/green.h
+++ b/psycopg/green.h
@@ -59,7 +59,7 @@ HIDDEN PyObject *psyco_set_wait_callback(PyObject *self, PyObject *obj);
 HIDDEN PyObject *psyco_get_wait_callback(PyObject *self, PyObject *obj);
 
 HIDDEN int psyco_green(void);
-HIDDEN PyObject *psyco_wait(connectionObject *conn);
+HIDDEN int psyco_wait(connectionObject *conn);
 HIDDEN PGresult *psyco_exec_green(connectionObject *conn, const char *command);
 
 #ifdef __cplusplus