From 9f06df1820459bf9031c8cf092fa993839b08f46 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 3 Dec 2012 02:49:06 +0000 Subject: [PATCH 1/5] Fixed signature for METH_NOARGS functions --- psycopg/connection_type.c | 24 ++++++++++++------------ psycopg/cursor_type.c | 10 +++++----- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index 08037079..69ab7c83 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -124,7 +124,7 @@ exit: #define psyco_conn_close_doc "close() -- Close the connection." static PyObject * -psyco_conn_close(connectionObject *self, PyObject *args) +psyco_conn_close(connectionObject *self) { Dprintf("psyco_conn_close: closing connection at %p", self); conn_close(self); @@ -140,7 +140,7 @@ psyco_conn_close(connectionObject *self, PyObject *args) #define psyco_conn_commit_doc "commit() -- Commit all changes to database." static PyObject * -psyco_conn_commit(connectionObject *self, PyObject *args) +psyco_conn_commit(connectionObject *self) { EXC_IF_CONN_CLOSED(self); EXC_IF_CONN_ASYNC(self, commit); @@ -160,7 +160,7 @@ psyco_conn_commit(connectionObject *self, PyObject *args) "rollback() -- Roll back all changes done to database." static PyObject * -psyco_conn_rollback(connectionObject *self, PyObject *args) +psyco_conn_rollback(connectionObject *self) { EXC_IF_CONN_CLOSED(self); EXC_IF_CONN_ASYNC(self, rollback); @@ -234,7 +234,7 @@ exit: "tpc_prepare() -- perform the first phase of a two-phase transaction." static PyObject * -psyco_conn_tpc_prepare(connectionObject *self, PyObject *args) +psyco_conn_tpc_prepare(connectionObject *self) { EXC_IF_CONN_CLOSED(self); EXC_IF_CONN_ASYNC(self, tpc_prepare); @@ -378,7 +378,7 @@ psyco_conn_tpc_rollback(connectionObject *self, PyObject *args) "tpc_recover() -- returns a list of pending transaction IDs." static PyObject * -psyco_conn_tpc_recover(connectionObject *self, PyObject *args) +psyco_conn_tpc_recover(connectionObject *self) { EXC_IF_CONN_CLOSED(self); EXC_IF_CONN_ASYNC(self, tpc_recover); @@ -652,7 +652,7 @@ psyco_conn_set_client_encoding(connectionObject *self, PyObject *args) "get_transaction_status() -- Get backend transaction status." static PyObject * -psyco_conn_get_transaction_status(connectionObject *self, PyObject *args) +psyco_conn_get_transaction_status(connectionObject *self) { EXC_IF_CONN_CLOSED(self); @@ -756,7 +756,7 @@ psyco_conn_lobject(connectionObject *self, PyObject *args, PyObject *keywds) "get_backend_pid() -- Get backend process id." static PyObject * -psyco_conn_get_backend_pid(connectionObject *self, PyObject *args) +psyco_conn_get_backend_pid(connectionObject *self) { EXC_IF_CONN_CLOSED(self); @@ -769,7 +769,7 @@ psyco_conn_get_backend_pid(connectionObject *self, PyObject *args) "reset() -- Reset current connection to defaults." static PyObject * -psyco_conn_reset(connectionObject *self, PyObject *args) +psyco_conn_reset(connectionObject *self) { int res; @@ -797,7 +797,7 @@ psyco_conn_get_exception(PyObject *self, void *closure) } static PyObject * -psyco_conn_poll(connectionObject *self, PyObject *args) +psyco_conn_poll(connectionObject *self) { int res; @@ -819,7 +819,7 @@ psyco_conn_poll(connectionObject *self, PyObject *args) "fileno() -> int -- Return file descriptor associated to database connection." static PyObject * -psyco_conn_fileno(connectionObject *self, PyObject *args) +psyco_conn_fileno(connectionObject *self) { long int socket; @@ -838,7 +838,7 @@ psyco_conn_fileno(connectionObject *self, PyObject *args) "executing an asynchronous operation." static PyObject * -psyco_conn_isexecuting(connectionObject *self, PyObject *args) +psyco_conn_isexecuting(connectionObject *self) { /* synchronous connections will always return False */ if (self->async == 0) { @@ -870,7 +870,7 @@ psyco_conn_isexecuting(connectionObject *self, PyObject *args) "cancel() -- cancel the current operation" static PyObject * -psyco_conn_cancel(connectionObject *self, PyObject *args) +psyco_conn_cancel(connectionObject *self) { char errbuf[256]; diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index bc8195ab..b2c55aa3 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -50,7 +50,7 @@ extern PyObject *pyPsycopgTzFixedOffsetTimezone; "close() -- Close the cursor." static PyObject * -psyco_curs_close(cursorObject *self, PyObject *args) +psyco_curs_close(cursorObject *self) { EXC_IF_ASYNC_IN_PROGRESS(self, close); @@ -761,7 +761,7 @@ exit: } static PyObject * -psyco_curs_fetchone(cursorObject *self, PyObject *args) +psyco_curs_fetchone(cursorObject *self) { PyObject *res; @@ -953,7 +953,7 @@ exit: "Return `!None` when no more data is available.\n" static PyObject * -psyco_curs_fetchall(cursorObject *self, PyObject *args) +psyco_curs_fetchall(cursorObject *self) { int i, size; PyObject *list = NULL; @@ -1085,7 +1085,7 @@ exit: "sets) and will raise a NotSupportedError exception." static PyObject * -psyco_curs_nextset(cursorObject *self, PyObject *args) +psyco_curs_nextset(cursorObject *self) { EXC_IF_CURS_CLOSED(self); @@ -1674,7 +1674,7 @@ cursor_next(PyObject *self) if (NULL == ((cursorObject*)self)->name) { /* we don't parse arguments: psyco_curs_fetchone will do that for us */ - res = psyco_curs_fetchone((cursorObject*)self, NULL); + res = psyco_curs_fetchone((cursorObject*)self); /* convert a None to NULL to signal the end of iteration */ if (res && res == Py_None) { From cc605032f5c0c00bee21ab34e5e13b9d866a795e Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 3 Dec 2012 02:50:24 +0000 Subject: [PATCH 2/5] Added support for with statement for connection and cursor The implementation should be conform to the DBAPI, although the "with" extension has not been released yet. --- psycopg/connection_type.c | 48 ++++++++++++ psycopg/cursor_type.c | 38 +++++++++ tests/__init__.py | 7 ++ tests/test_with.py | 157 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 250 insertions(+) create mode 100755 tests/test_with.py diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index 69ab7c83..c1d6176e 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -389,6 +389,50 @@ psyco_conn_tpc_recover(connectionObject *self) } +#define psyco_conn_enter_doc \ +"__enter__ -> self" + +static PyObject * +psyco_conn_enter(connectionObject *self) +{ + EXC_IF_CONN_CLOSED(self); + + Py_INCREF(self); + return (PyObject *)self; +} + + +#define psyco_conn_exit_doc \ +"__exit__ -- commit if no exception, else roll back" + +static PyObject * +psyco_conn_exit(connectionObject *self, PyObject *args) +{ + PyObject *type, *name, *tb; + PyObject *tmp = NULL; + PyObject *rv = NULL; + + if (!PyArg_ParseTuple(args, "OOO", &type, &name, &tb)) { + goto exit; + } + + if (type == Py_None) { + if (!(tmp = psyco_conn_commit(self))) { goto exit; } + } else { + if (!(tmp = psyco_conn_rollback(self))) { goto exit; } + } + + /* success (of the commit or rollback, there may have been an exception in + * the block). Return None to avoid swallowing the exception */ + rv = Py_None; + Py_INCREF(rv); + +exit: + Py_XDECREF(tmp); + return rv; +} + + #ifdef PSYCOPG_EXTENSIONS @@ -924,6 +968,10 @@ static struct PyMethodDef connectionObject_methods[] = { METH_VARARGS, psyco_conn_tpc_rollback_doc}, {"tpc_recover", (PyCFunction)psyco_conn_tpc_recover, METH_NOARGS, psyco_conn_tpc_recover_doc}, + {"__enter__", (PyCFunction)psyco_conn_enter, + METH_NOARGS, psyco_conn_enter_doc}, + {"__exit__", (PyCFunction)psyco_conn_exit, + METH_VARARGS, psyco_conn_exit_doc}, #ifdef PSYCOPG_EXTENSIONS {"set_session", (PyCFunction)psyco_conn_set_session, METH_VARARGS|METH_KEYWORDS, psyco_conn_set_session_doc}, diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index b2c55aa3..5e17bff1 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -1201,6 +1201,40 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs) } +#define psyco_curs_enter_doc \ +"__enter__ -> self" + +static PyObject * +psyco_curs_enter(cursorObject *self) +{ + Py_INCREF(self); + return (PyObject *)self; +} + +#define psyco_curs_exit_doc \ +"__exit__ -- close the cursor" + +static PyObject * +psyco_curs_exit(cursorObject *self, PyObject *args) +{ + PyObject *tmp = NULL; + PyObject *rv = NULL; + + /* don't care about the arguments here: don't need to parse them */ + + if (!(tmp = psyco_curs_close(self))) { goto exit; } + + /* success (of curs.close()). + * Return None to avoid swallowing the exception */ + rv = Py_None; + Py_INCREF(rv); + +exit: + Py_XDECREF(tmp); + return rv; +} + + #ifdef PSYCOPG_EXTENSIONS /* Return a newly allocated buffer containing the list of columns to be @@ -1716,6 +1750,10 @@ static struct PyMethodDef cursorObject_methods[] = { /* DBAPI-2.0 extensions */ {"scroll", (PyCFunction)psyco_curs_scroll, METH_VARARGS|METH_KEYWORDS, psyco_curs_scroll_doc}, + {"__enter__", (PyCFunction)psyco_curs_enter, + METH_NOARGS, psyco_curs_enter_doc}, + {"__exit__", (PyCFunction)psyco_curs_exit, + METH_VARARGS, psyco_curs_exit_doc}, /* psycopg extensions */ #ifdef PSYCOPG_EXTENSIONS {"cast", (PyCFunction)psyco_curs_cast, diff --git a/tests/__init__.py b/tests/__init__.py index df8e8cd1..3e677d85 100755 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -45,6 +45,11 @@ import test_transaction import test_types_basic import test_types_extras +if sys.version_info[:2] >= (2, 5): + import test_with +else: + test_with = None + def test_suite(): # If connection to test db fails, bail out early. import psycopg2 @@ -76,6 +81,8 @@ def test_suite(): suite.addTest(test_transaction.test_suite()) suite.addTest(test_types_basic.test_suite()) suite.addTest(test_types_extras.test_suite()) + if test_with: + suite.addTest(test_with.test_suite()) return suite if __name__ == '__main__': diff --git a/tests/test_with.py b/tests/test_with.py new file mode 100755 index 00000000..51889270 --- /dev/null +++ b/tests/test_with.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python + +# test_ctxman.py - unit test for connection and cursor used as context manager +# +# Copyright (C) 2012 Daniele Varrazzo +# +# psycopg2 is free software: you can redistribute it and/or modify it +# under the terms of the GNU Lesser General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# In addition, as a special exception, the copyright holders give +# permission to link this program with the OpenSSL library (or with +# modified versions of OpenSSL that use the same license as OpenSSL), +# and distribute linked combinations including the two. +# +# You must obey the GNU Lesser General Public License in all respects for +# all of the code used other than OpenSSL. +# +# psycopg2 is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public +# License for more details. + + +from __future__ import with_statement + +import psycopg2 +import psycopg2.extensions as ext + +from testconfig import dsn +from testutils import unittest + +class TestMixin(object): + def setUp(self): + self.conn = conn = psycopg2.connect(dsn) + curs = conn.cursor() + try: + curs.execute("delete from test_with") + conn.commit() + except psycopg2.ProgrammingError: + # assume table doesn't exist + conn.rollback() + curs.execute("create table test_with (id integer primary key)") + conn.commit() + + def tearDown(self): + self.conn.close() + + +class WithConnectionTestCase(TestMixin, unittest.TestCase): + def test_with_ok(self): + with self.conn as conn: + self.assert_(self.conn is conn) + self.assertEqual(conn.status, ext.STATUS_READY) + curs = conn.cursor() + curs.execute("insert into test_with values (1)") + self.assertEqual(conn.status, ext.STATUS_BEGIN) + + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), [(1,)]) + + def test_with_connect_idiom(self): + with psycopg2.connect(dsn) as conn: + self.assertEqual(conn.status, ext.STATUS_READY) + curs = conn.cursor() + curs.execute("insert into test_with values (2)") + self.assertEqual(conn.status, ext.STATUS_BEGIN) + + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), [(2,)]) + + def test_with_error_db(self): + def f(): + with self.conn as conn: + curs = conn.cursor() + curs.execute("insert into test_with values ('a')") + + self.assertRaises(psycopg2.DataError, f) + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), []) + + def test_with_error_python(self): + def f(): + with self.conn as conn: + curs = conn.cursor() + curs.execute("insert into test_with values (3)") + 1/0 + + self.assertRaises(ZeroDivisionError, f) + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), []) + + def test_with_closed(self): + def f(): + with self.conn: + pass + + self.conn.close() + self.assertRaises(psycopg2.InterfaceError, f) + + +class WithCursorTestCase(TestMixin, unittest.TestCase): + def test_with_ok(self): + with self.conn as conn: + with conn.cursor() as curs: + curs.execute("insert into test_with values (4)") + self.assert_(not curs.closed) + self.assertEqual(self.conn.status, ext.STATUS_BEGIN) + self.assert_(curs.closed) + + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), [(4,)]) + + def test_with_error(self): + try: + with self.conn as conn: + with conn.cursor() as curs: + curs.execute("insert into test_with values (5)") + 1/0 + except ZeroDivisionError: + pass + + self.assertEqual(self.conn.status, ext.STATUS_READY) + self.assert_(not self.conn.closed) + self.assert_(curs.closed) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), []) + + +def test_suite(): + return unittest.TestLoader().loadTestsFromName(__name__) + +if __name__ == "__main__": + unittest.main() From c2f284cd3b7d3b33f4de348aaac62590feda8c78 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 3 Dec 2012 03:18:51 +0000 Subject: [PATCH 3/5] Added documentation for the with statement --- doc/src/connection.rst | 8 ++++++++ doc/src/cursor.rst | 5 +++++ doc/src/usage.rst | 24 ++++++++++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/doc/src/connection.rst b/doc/src/connection.rst index 997f96e9..f7ff4b19 100644 --- a/doc/src/connection.rst +++ b/doc/src/connection.rst @@ -74,6 +74,10 @@ The ``connection`` class automatically open, commands have immediate effect. See :ref:`transactions-control` for details. + .. versionchanged:: 2.5 if the connection is used in a ``with`` + statement, the method is automatically called if no exception is + raised in the ``with`` block. + .. index:: pair: Transaction; Rollback @@ -84,6 +88,10 @@ The ``connection`` class connection without committing the changes first will cause an implicit rollback to be performed. + .. versionchanged:: 2.5 if the connection is used in a ``with`` + statement, the method is automatically called if an exception is + raised in the ``with`` block. + .. method:: close() diff --git a/doc/src/cursor.rst b/doc/src/cursor.rst index 204fce21..62be5e3c 100644 --- a/doc/src/cursor.rst +++ b/doc/src/cursor.rst @@ -83,6 +83,11 @@ The ``cursor`` class The cursor will be unusable from this point forward; an `~psycopg2.InterfaceError` will be raised if any operation is attempted with the cursor. + + .. versionchanged:: 2.5 if the cursor is used in a ``with`` statement, + the method is automatically called at the end of the ``with`` + block. + .. attribute:: closed diff --git a/doc/src/usage.rst b/doc/src/usage.rst index a2c9b3fe..2586d7d8 100644 --- a/doc/src/usage.rst +++ b/doc/src/usage.rst @@ -548,6 +548,30 @@ change the isolation level. See the `~connection.set_session()` method for all the details. +.. index:: + single: with statement + +``with`` statement +^^^^^^^^^^^^^^^^^^ + +Starting from version 2.5, psycopg2's connections and cursors are *context +managers* and can be used with the ``with`` statement:: + + with psycopg2.connect(DSN) as conn: + with conn.cursor() as curs: + curs.execute(SQL) + +When a connection exits the ``with`` block, if no exception has been raised by +the block, the transaction is committed. In case of exception the transaction +is rolled back. In no case the connection is closed: a connection can be used +in more than a ``with`` statement and each ``with`` block is effectively +wrapped in a transaction. + +When a cursor exits the ``with`` block it is closed, releasing any resource +eventually associated with it. The state of the transaction is not affected. + + + .. index:: pair: Server side; Cursor pair: Named; Cursor From 12645db754f9f665f8959a6072628ddb094f9abf Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 3 Dec 2012 03:37:47 +0000 Subject: [PATCH 4/5] Make sure to call subclasses methods on context exit --- psycopg/connection_type.c | 8 ++++-- psycopg/cursor_type.c | 4 ++- tests/test_with.py | 55 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index c1d6176e..b5fd0789 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -417,9 +417,13 @@ psyco_conn_exit(connectionObject *self, PyObject *args) } if (type == Py_None) { - if (!(tmp = psyco_conn_commit(self))) { goto exit; } + if (!(tmp = PyObject_CallMethod((PyObject *)self, "commit", ""))) { + goto exit; + } } else { - if (!(tmp = psyco_conn_rollback(self))) { goto exit; } + if (!(tmp = PyObject_CallMethod((PyObject *)self, "rollback", ""))) { + goto exit; + } } /* success (of the commit or rollback, there may have been an exception in diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index 5e17bff1..9570f914 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -1222,7 +1222,9 @@ psyco_curs_exit(cursorObject *self, PyObject *args) /* don't care about the arguments here: don't need to parse them */ - if (!(tmp = psyco_curs_close(self))) { goto exit; } + if (!(tmp = PyObject_CallMethod((PyObject *)self, "close", ""))) { + goto exit; + } /* success (of curs.close()). * Return None to avoid swallowing the exception */ diff --git a/tests/test_with.py b/tests/test_with.py index 51889270..f43e6db4 100755 --- a/tests/test_with.py +++ b/tests/test_with.py @@ -115,6 +115,48 @@ class WithConnectionTestCase(TestMixin, unittest.TestCase): self.conn.close() self.assertRaises(psycopg2.InterfaceError, f) + def test_subclass_commit(self): + commits = [] + class MyConn(ext.connection): + def commit(self): + commits.append(None) + super(MyConn, self).commit() + + with psycopg2.connect(dsn, connection_factory=MyConn) as conn: + curs = conn.cursor() + curs.execute("insert into test_with values (10)") + + self.assertEqual(conn.status, ext.STATUS_READY) + self.assert_(commits) + + curs = self.conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), [(10,)]) + + def test_subclass_rollback(self): + rollbacks = [] + class MyConn(ext.connection): + def rollback(self): + rollbacks.append(None) + super(MyConn, self).rollback() + + try: + with psycopg2.connect(dsn, connection_factory=MyConn) as conn: + curs = conn.cursor() + curs.execute("insert into test_with values (11)") + 1/0 + except ZeroDivisionError: + pass + else: + self.assert_("exception not raised") + + self.assertEqual(conn.status, ext.STATUS_READY) + self.assert_(rollbacks) + + curs = conn.cursor() + curs.execute("select * from test_with") + self.assertEqual(curs.fetchall(), []) + class WithCursorTestCase(TestMixin, unittest.TestCase): def test_with_ok(self): @@ -149,6 +191,19 @@ class WithCursorTestCase(TestMixin, unittest.TestCase): curs.execute("select * from test_with") self.assertEqual(curs.fetchall(), []) + def test_subclass(self): + closes = [] + class MyCurs(ext.cursor): + def close(self): + closes.append(None) + super(MyCurs, self).close() + + with self.conn.cursor(cursor_factory=MyCurs) as curs: + self.assert_(isinstance(curs, MyCurs)) + + self.assert_(curs.closed) + self.assert_(closes) + def test_suite(): return unittest.TestLoader().loadTestsFromName(__name__) From ec34b9bed6e9b1d3f3638165fb74e801b37bede0 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Tue, 4 Dec 2012 00:26:48 +0000 Subject: [PATCH 5/5] Mention context managers in NEWS file --- NEWS | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NEWS b/NEWS index 8f4f258c..6ceb6fb5 100644 --- a/NEWS +++ b/NEWS @@ -3,6 +3,8 @@ What's new in psycopg 2.5 - Added JSON adaptation. - Added support for PostgreSQL 9.2 range types. + - 'connection' and 'cursor' objects can be used in 'with' statements + as context managers as specified by recent DBAPI extension. - Added support for backward scrollable cursors. Thanks to Jon Nelson for the initial patch (ticket #108). - Added a simple way to customize casting of composite types into Python