Handle initial connection queries being sent partly

The CONN_STATUS_SENT_* statuses were not being handled at all, and
they indicate that a query has been sent, but not fully, so the client
should wait for the socket to become writable again and flush the output.

2 files changed, 17 insertions, 7 deletions

diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 0fce84d..baca47f 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -387,7 +387,7 @@ conn_connect(connectionObject *self, long int async)
 PyObject *
 conn_poll_send(connectionObject *self)
 {
-    const char *query;
+    const char *query = NULL;
     int next_status;
     int ret;
 
@@ -404,6 +404,12 @@ conn_poll_send(connectionObject *self)
         query = psyco_client_encoding;
         next_status = CONN_STATUS_SENT_CLIENT_ENCODING;
         break;
+    case CONN_STATUS_SENT_DATESTYLE:
+    case CONN_STATUS_SENT_CLIENT_ENCODING:
+        /* the query has only been partially sent */
+        query = NULL;
+        next_status = self->status;
+        break;
     default:
         /* unexpected state, error out */
         PyErr_Format(OperationalError,
@@ -416,12 +422,14 @@ conn_poll_send(connectionObject *self)
     Py_BEGIN_ALLOW_THREADS;
     pthread_mutex_lock(&(self->lock));
 
-    if (PQsendQuery(self->pgconn, query) != 1) {
-        pthread_mutex_unlock(&(self->lock));
-        Py_BLOCK_THREADS;
-        PyErr_SetString(OperationalError,
-                        PQerrorMessage(self->pgconn));
-        return NULL;
+    if (query != NULL) {
+        if (PQsendQuery(self->pgconn, query) != 1) {
+            pthread_mutex_unlock(&(self->lock));
+            Py_BLOCK_THREADS;
+            PyErr_SetString(OperationalError,
+                            PQerrorMessage(self->pgconn));
+            return NULL;
+        }
     }
 
     if (PQflush(self->pgconn) == 0) {
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index f80eba1..efbe9e6 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -426,7 +426,9 @@ psyco_conn_poll(connectionObject *self)
     switch (self->status) {
 
     case CONN_STATUS_SEND_DATESTYLE:
+    case CONN_STATUS_SENT_DATESTYLE:
     case CONN_STATUS_SEND_CLIENT_ENCODING:
+    case CONN_STATUS_SENT_CLIENT_ENCODING:
         /* these mean that we need to wait for the socket to become writable
            to send the rest of our query */
         return conn_poll_send(self);