From 389ad08965ff10f22249f1c3d5693fde121ff9be Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 2 Apr 2010 02:00:48 +0100
Subject: [PATCH] Use the wait callback during connection if set.

---
 psycopg/connection_int.c  | 56 +++++++++++++++++++++++++++++++--------
 psycopg/connection_type.c |  2 +-
 2 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 5183bbef..049efa0b 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -33,6 +33,7 @@
 #include "psycopg/connection.h"
 #include "psycopg/cursor.h"
 #include "psycopg/pqpath.h"
+#include "psycopg/green.h"
 
 /* conn_notice_callback - process notices */
 
@@ -319,12 +320,24 @@ int
 conn_sync_connect(connectionObject *self)
 {
     PGconn *pgconn;
+    PyObject *wait_rv;
+    int green;
 
-    Py_BEGIN_ALLOW_THREADS;
-    self->pgconn = pgconn = PQconnectdb(self->dsn);
-    Py_END_ALLOW_THREADS;
-
-    Dprintf("conn_connect: new postgresql connection at %p", pgconn);
+    /* store this value to prevent inconsistencies due to a change
+     * in the middle of the function. */
+    green = psyco_green();
+    if (!green) {
+        Py_BEGIN_ALLOW_THREADS;
+        self->pgconn = pgconn = PQconnectdb(self->dsn);
+        Py_END_ALLOW_THREADS;
+        Dprintf("conn_connect: new postgresql connection at %p", pgconn);
+    }
+    else {
+        Py_BEGIN_ALLOW_THREADS;
+        self->pgconn = pgconn = PQconnectStart(self->dsn);
+        Py_END_ALLOW_THREADS;
+        Dprintf("conn_connect: new green postgresql connection at %p", pgconn);
+    }
 
     if (pgconn == NULL)
     {
@@ -341,18 +354,39 @@ conn_sync_connect(connectionObject *self)
 
     PQsetNoticeProcessor(pgconn, conn_notice_callback, (void*)self);
 
-    if (conn_setup(self, pgconn) == -1)
+#ifdef HAVE_PQPROTOCOL3
+    self->protocol = PQprotocolVersion(pgconn);
+#else
+    self->protocol = 2;
+#endif
+
+    Dprintf("conn_connect: using protocol %d", self->protocol);
+
+    self->server_version = (int)PQserverVersion(pgconn);
+
+    /* if the connection is green, wait to finish connection */
+    if (green) {
+        wait_rv = psyco_wait((PyObject *)self, Py_None);
+        if (wait_rv) {
+            Py_DECREF(wait_rv);
+        } else {
+            return -1;
+        }
+    }
+
+    /* From here the connection is considered ready: with the new status,
+     * poll() will use PQisBusy instead of PQconnectPoll.
+     */
+    self->status = CONN_STATUS_READY;
+
+    if (conn_setup(self, self->pgconn) == -1) {
         return -1;
+    }
 
     if (pq_set_non_blocking(self, 1, 1) != 0) {
         return -1;
     }
 
-    self->protocol = conn_get_protocol_version(pgconn);
-    Dprintf("conn_connect: using protocol %d", self->protocol);
-
-    self->server_version = (int)PQserverVersion(pgconn);
-
     return 0;
 }
 
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 42ada066..54633051 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -702,7 +702,7 @@ connection_setup(connectionObject *self, const char *dsn, long int async)
     self->notifies = PyList_New(0);
     self->closed = 0;
     self->async = async;
-    self->status = async ? CONN_STATUS_SETUP : CONN_STATUS_READY;
+    self->status = CONN_STATUS_SETUP;
     self->critical = NULL;
     self->async_cursor = NULL;
     self->async_status = ASYNC_DONE;