From 04d66a6c828dc69415d176714213d43038de801a Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 2 Apr 2010 01:59:31 +0100
Subject: [PATCH] Added poll implementation for sync connection with external
 wait.

---
 psycopg/connection.h      |  2 ++
 psycopg/connection_int.c  | 41 +++++++++++++++++++++++++++++++++++++++
 psycopg/connection_type.c | 21 ++++++++++++++++----
 3 files changed, 60 insertions(+), 4 deletions(-)

diff --git a/psycopg/connection.h b/psycopg/connection.h
index 3a57e8da..8eae85ac 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -59,6 +59,7 @@ extern "C" {
 #define PSYCO_POLL_OK    0
 #define PSYCO_POLL_READ  1
 #define PSYCO_POLL_WRITE 2
+#define PSYCO_POLL_ERROR 3
 
 /* Hard limit on the notices stored by the Python connection */
 #define CONN_NOTICES_LIMIT 50
@@ -134,6 +135,7 @@ HIDDEN int  conn_set_client_encoding(connectionObject *self, const char *enc);
 HIDDEN PyObject *conn_poll_send(connectionObject *self);
 HIDDEN PyObject *conn_poll_fetch(connectionObject *self);
 HIDDEN PyObject *conn_poll_ready(connectionObject *self);
+HIDDEN PyObject *conn_poll_green(connectionObject *self);
 
 /* exception-raising macros */
 #define EXC_IF_CONN_CLOSED(self) if ((self)->closed > 0) { \
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 66a22524..5183bbef 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -629,6 +629,47 @@ conn_poll_ready(connectionObject *self)
     }
 }
 
+/* conn_poll_green - poll a *sync* connection with external wait */
+
+PyObject *
+conn_poll_green(connectionObject *self)
+{
+    int res = PSYCO_POLL_ERROR;
+
+    switch (self->status) {
+    case CONN_STATUS_SETUP:
+        Dprintf("conn_poll: status = CONN_STATUS_SETUP");
+        self->status = CONN_STATUS_ASYNC;
+        res =  PSYCO_POLL_WRITE;
+        break;
+
+    case CONN_STATUS_ASYNC:
+        Dprintf("conn_poll: status = CONN_STATUS_ASYNC");
+        switch (PQconnectPoll(self->pgconn)) {
+        case PGRES_POLLING_OK:
+            res = PSYCO_POLL_OK;
+            break;
+        case PGRES_POLLING_READING:
+            res = PSYCO_POLL_READ;
+            break;
+        case PGRES_POLLING_WRITING:
+            res = PSYCO_POLL_WRITE;
+            break;
+        case PGRES_POLLING_FAILED:
+        case PGRES_POLLING_ACTIVE:
+            res = PSYCO_POLL_ERROR;
+            break;
+        }
+        break;
+
+    default:
+        Dprintf("conn_poll: in unexpected state");
+        res = PSYCO_POLL_ERROR;
+    }
+
+    return PyInt_FromLong(res);
+}
+
 /* conn_close - do anything needed to shut down the connection */
 
 void
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 84ce6ed5..42ada066 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -399,8 +399,6 @@ psyco_conn_reset(connectionObject *self)
     return Py_None;
 }
 
-#endif
-
 static PyObject *
 psyco_conn_get_exception(PyObject *self, void *closure)
 {
@@ -411,13 +409,13 @@ psyco_conn_get_exception(PyObject *self, void *closure)
 }
 
 #define psyco_conn_poll_doc \
-"poll() -- return POLL_OK if the connection has been estabilished, " \
+"poll() -- return POLL_OK if the operation has finished, "           \
  "POLL_READ if the application should be waiting "                   \
  "for the socket to be readable or POLL_WRITE "                      \
  "if the socket should be writable."
 
 static PyObject *
-psyco_conn_poll(connectionObject *self)
+psyco_conn_poll_async(connectionObject *self)
 {
     PostgresPollingStatusType poll_status;
 
@@ -524,6 +522,18 @@ psyco_conn_poll(connectionObject *self)
     return PyInt_FromLong(PSYCO_POLL_WRITE);
 }
 
+static PyObject *
+psyco_conn_poll(connectionObject *self)
+{
+    EXC_IF_CONN_CLOSED(self);
+
+    if (self->async) {
+        return psyco_conn_poll_async(self);
+    } else {
+        return conn_poll_green(self);
+    }
+}
+
 
 /* extension: fileno - return the file descriptor of the connection */
 
@@ -579,6 +589,9 @@ psyco_conn_isexecuting(connectionObject *self)
     return Py_False;
 }
 
+#endif  /* PSYCOPG_EXTENSIONS */
+
+
 /** the connection object **/