Implemented connection.reset()

ChangeLog

@@ -1,5 +1,10 @@
 2009-08-08  Federico Di Gregorio  <fog@initd.org>
 
+	* Implemented connection.reset() method to reset the connection to
+	well-know default parameters. This is much faster than closing and
+	reopening the connection. (Suggested by a bug report by Glenn
+	Maynard.)
+
 	* psycopg/cursor_type.c: unified size macro definitions in COPY TO
 	and COPY FROM operations: now the buffer for column names is 8192
 	bytes that should be enough even for very large tables.

psycopg/connection.h

@@ -87,6 +87,7 @@ typedef struct {
 /* C-callable functions in connection_int.c and connection_ext.c */
 HIDDEN void conn_notice_process(connectionObject *self);
 HIDDEN void conn_notice_clean(connectionObject *self);
+HIDDEN int  conn_setup(connectionObject *self, PGconn *pgconn);
 HIDDEN int  conn_connect(connectionObject *self);
 HIDDEN void conn_close(connectionObject *self);
 HIDDEN int  conn_commit(connectionObject *self);

psycopg/connection_int.c

@@ -107,12 +107,11 @@ conn_notice_clean(connectionObject *self)
     pthread_mutex_unlock(&self->lock);    
 }
 
-/* conn_connect - execute a connection to the database */
+/* conn_setup - setup and read basic information about the connection */
 
 int
-conn_connect(connectionObject *self)
+conn_setup(connectionObject *self, PGconn *pgconn)
 {
-    PGconn *pgconn;
     PGresult *pgres;
     const char *data, *tmp;
     const char *scs;    /* standard-conforming strings */
@@ -129,27 +128,9 @@ conn_connect(connectionObject *self)
     static const char lvl2a[] = "repeatable read";
     static const char lvl2b[] = "serializable";
 
-    Py_BEGIN_ALLOW_THREADS;
-    pgconn = PQconnectdb(self->dsn);
-    Py_END_ALLOW_THREADS;
-
-    Dprintf("conn_connect: new postgresql connection at %p", pgconn);
-
-    if (pgconn == NULL)
-    {
-        Dprintf("conn_connect: PQconnectdb(%s) FAILED", self->dsn);
-        PyErr_SetString(OperationalError, "PQconnectdb() failed");
-        return -1;
-    }
-    else if (PQstatus(pgconn) == CONNECTION_BAD)
-    {
-        Dprintf("conn_connect: PQconnectdb(%s) returned BAD", self->dsn);
-        PyErr_SetString(OperationalError, PQerrorMessage(pgconn));
-        PQfinish(pgconn);
-        return -1;
-    }
-
-    PQsetNoticeProcessor(pgconn, conn_notice_callback, (void*)self);
+    if (self->encoding) free(self->encoding);
+    self->equote = 0;
+    self->isolation_level = 0;
 
     /*
      * The presence of the 'standard_conforming_strings' parameter
@@ -236,6 +217,41 @@ conn_connect(connectionObject *self)
         self->isolation_level = 2;
     CLEARPGRES(pgres);
 
+    return 0;
+}
+
+/* conn_connect - execute a connection to the database */
+
+int
+conn_connect(connectionObject *self)
+{
+    PGconn *pgconn;
+
+    Py_BEGIN_ALLOW_THREADS;
+    pgconn = PQconnectdb(self->dsn);
+    Py_END_ALLOW_THREADS;
+
+    Dprintf("conn_connect: new postgresql connection at %p", pgconn);
+
+    if (pgconn == NULL)
+    {
+        Dprintf("conn_connect: PQconnectdb(%s) FAILED", self->dsn);
+        PyErr_SetString(OperationalError, "PQconnectdb() failed");
+        return -1;
+    }
+    else if (PQstatus(pgconn) == CONNECTION_BAD)
+    {
+        Dprintf("conn_connect: PQconnectdb(%s) returned BAD", self->dsn);
+        PyErr_SetString(OperationalError, PQerrorMessage(pgconn));
+        PQfinish(pgconn);
+        return -1;
+    }
+
+    PQsetNoticeProcessor(pgconn, conn_notice_callback, (void*)self);
+
+    if (conn_setup(self, pgconn) == -1)
+        return -1;
+
     if (PQsetnonblocking(pgconn, 1) != 0) {
         Dprintf("conn_connect: PQsetnonblocking() FAILED");
         PyErr_SetString(OperationalError, "PQsetnonblocking() failed");

psycopg/connection_type.c

@@ -351,6 +351,31 @@ psyco_conn_get_backend_pid(connectionObject *self)
     return PyInt_FromLong((long)PQbackendPID(self->pgconn));
 }
 
+/* reset the currect connection */
+
+#define psyco_conn_reset_doc \
+"reset() -- Reset current connection to defaults."
+
+static PyObject *
+psyco_conn_reset(connectionObject *self)
+{
+    int res;
+
+    EXC_IF_CONN_CLOSED(self);
+
+    if (pq_reset(self) < 0)
+        return NULL;
+
+    pthread_mutex_lock(&self->lock);
+    res = conn_setup(self, self->pgconn);
+    pthread_mutex_unlock(&self->lock);
+    if (res < 0)
+        return NULL;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 #endif
 
 static PyObject *
@@ -389,6 +414,8 @@ static struct PyMethodDef connectionObject_methods[] = {
      METH_NOARGS, psyco_conn_get_backend_pid_doc},
     {"lobject", (PyCFunction)psyco_conn_lobject,
      METH_VARARGS|METH_KEYWORDS, psyco_conn_lobject_doc},
+    {"reset", (PyCFunction)psyco_conn_reset,
+     METH_NOARGS, psyco_conn_reset_doc},
 #endif
     {NULL}
 };
@@ -469,6 +496,7 @@ connection_setup(connectionObject *self, const char *dsn)
     self->string_types = PyDict_New();
     self->binary_types = PyDict_New();
     self->notice_pending = NULL;
+    self->encoding = NULL;
 
     pthread_mutex_init(&(self->lock), NULL);
 

psycopg/pqpath.c

@@ -401,7 +401,8 @@ pq_begin_locked(connectionObject *conn, PGresult **pgres, char **error)
 /* pq_commit - send an END, if necessary
 
    This function should be called while holding the global interpreter
-   lock. */
+   lock.
+*/
 
 int
 pq_commit(connectionObject *conn)
@@ -498,6 +499,69 @@ pq_abort(connectionObject *conn)
     return retvalue;
 }
 
+/* pq_reset - reset the connection
+
+   This function should be called while holding the global interpreter
+   lock.
+
+   The _locked version of this function should be called on a locked
+   connection without holding the global interpreter lock.
+*/
+
+int
+pq_reset_locked(connectionObject *conn, PGresult **pgres, char **error)
+{
+    int retvalue = -1;
+
+    Dprintf("pq_reset_locked: pgconn = %p, isolevel = %ld, status = %d",
+            conn->pgconn, conn->isolation_level, conn->status);
+
+    conn->mark += 1;
+    pq_clear_async(conn);
+
+    if (conn->isolation_level > 0 && conn->status == CONN_STATUS_BEGIN) {
+        retvalue = pq_execute_command_locked(conn, "ABORT", pgres, error);
+        if (retvalue != 0) return retvalue;
+    }
+
+    retvalue = pq_execute_command_locked(conn, "RESET ALL", pgres, error);
+    if (retvalue != 0) return retvalue;
+
+    retvalue = pq_execute_command_locked(conn,
+        "SET SESSION AUTHORIZATION DEFAULT", pgres, error);
+    if (retvalue != 0) return retvalue;
+
+    conn->status = CONN_STATUS_READY;
+
+    return retvalue;
+}
+
+int
+pq_reset(connectionObject *conn)
+{
+    int retvalue = -1;
+    PGresult *pgres = NULL;
+    char *error = NULL;
+
+    Dprintf("pq_reset: pgconn = %p, isolevel = %ld, status = %d",
+            conn->pgconn, conn->isolation_level, conn->status);
+
+    Py_BEGIN_ALLOW_THREADS;
+    pthread_mutex_lock(&conn->lock);
+
+    retvalue = pq_reset_locked(conn, &pgres, &error);
+
+    pthread_mutex_unlock(&conn->lock);
+    Py_END_ALLOW_THREADS;
+
+    conn_notice_process(conn);
+
+    if (retvalue < 0)
+        pq_complete_error(conn, &pgres, &error);
+
+    return retvalue;
+}
+
 /* pq_is_busy - consume input and return connection status
 
    a status of 1 means that a call to pq_fetch will block, while a status of 0

psycopg/pqpath.h

@@ -39,6 +39,7 @@ HIDDEN int pq_commit(connectionObject *conn);
 HIDDEN int pq_abort_locked(connectionObject *conn, PGresult **pgres,
                            char **error);
 HIDDEN int pq_abort(connectionObject *conn);
+HIDDEN int pq_reset(connectionObject *conn);
 HIDDEN int pq_is_busy(connectionObject *conn);
 
 HIDDEN void pq_set_critical(connectionObject *conn, const char *msg);

tests/test_connection.py

@@ -28,7 +28,19 @@ class ConnectionTests(unittest.TestCase):
         conn.close()
         self.assertEqual(curs.closed, True)
 
+    def test_reset(self):
+        conn = self.connect()
+        # switch isolation level, then reset
+        level = conn.isolation_level
+        conn.set_isolation_level(0)
+        self.assertEqual(conn.isolation_level, 0)
+        conn.reset()
+        # now the isolation level should be equal to saved one
+        self.assertEqual(conn.isolation_level, level)
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)
 
+if __name__ == '__main__':
+    unittest.main(defaultTest='test_suite')