From 36aff2f73dd4eb1066c15b8a194752dd5e77b078 Mon Sep 17 00:00:00 2001 From: Federico Di Gregorio Date: Sun, 9 Aug 2009 16:19:08 +0200 Subject: [PATCH] Implemented connection.reset() --- ChangeLog | 5 +++ psycopg/connection.h | 1 + psycopg/connection_int.c | 66 ++++++++++++++++++++++-------------- psycopg/connection_type.c | 28 ++++++++++++++++ psycopg/pqpath.c | 70 +++++++++++++++++++++++++++++++++++++-- psycopg/pqpath.h | 1 + tests/test_connection.py | 12 +++++++ 7 files changed, 155 insertions(+), 28 deletions(-) diff --git a/ChangeLog b/ChangeLog index cf28058c..f2be2f2f 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,5 +1,10 @@ 2009-08-08 Federico Di Gregorio + * 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. diff --git a/psycopg/connection.h b/psycopg/connection.h index 701f46e0..5dce4270 100644 --- a/psycopg/connection.h +++ b/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); diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c index 8b54badc..cb41a758 100644 --- a/psycopg/connection_int.c +++ b/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"); @@ -249,7 +265,7 @@ conn_connect(connectionObject *self) self->protocol = 2; #endif Dprintf("conn_connect: using protocol %d", self->protocol); - + self->server_version = (int)PQserverVersion(pgconn); self->pgconn = pgconn; diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index 4bba61c1..4e6b42f0 100644 --- a/psycopg/connection_type.c +++ b/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); diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c index c687ef82..70bba7e2 100644 --- a/psycopg/pqpath.c +++ b/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) @@ -427,7 +428,7 @@ pq_commit(connectionObject *conn) pthread_mutex_unlock(&conn->lock); Py_END_ALLOW_THREADS; - + conn_notice_process(conn); if (retvalue < 0) @@ -489,7 +490,70 @@ pq_abort(connectionObject *conn) 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_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) diff --git a/psycopg/pqpath.h b/psycopg/pqpath.h index 8503c12f..2fd7652f 100644 --- a/psycopg/pqpath.h +++ b/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); diff --git a/tests/test_connection.py b/tests/test_connection.py index 306c7c36..5301ee37 100644 --- a/tests/test_connection.py +++ b/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')