'with' starts a transaction even on autocommit connections

Close #941

6 files changed, 127 insertions, 12 deletions

diff --git a/NEWS b/NEWS
index 47adaba..744d519 100644
--- a/NEWS
+++ b/NEWS
@@ -5,6 +5,8 @@ What's new in psycopg 2.9
 -------------------------
 
 - Dropped support for Python 2.7, 3.4, 3.5 (:tickets:`#1198, #1000, #1197`).
+- ``with connection`` starts a transaction on autocommit transactions too
+  (:ticket:`#941`).
 - Connection exceptions with sqlstate ``08XXX`` reclassified as
   `~psycopg2.OperationalError` (a subclass of the previously used
   `~psycopg2.DatabaseError`) (:ticket:`#1148`).
diff --git a/doc/src/usage.rst b/doc/src/usage.rst
index 0677a5b..3aafa90 100644
--- a/doc/src/usage.rst
+++ b/doc/src/usage.rst
@@ -832,6 +832,9 @@ and each ``with`` block is effectively wrapped in a separate transaction::
         finally:
             conn.close()
 
+.. versionchanged:: 2.9
+    ``with connection`` starts a transaction also on autocommit connections.
+
 
 .. index::
     pair: Server side; Cursor
diff --git a/psycopg/connection.h b/psycopg/connection.h
index a78cdce..e4a88ae 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -145,6 +145,9 @@ struct connectionObject {
 
     /* the pid this connection was created into */
     pid_t procpid;
+
+    /* inside a with block */
+    int entered;
 };
 
 /* map isolation level values into a numeric const */
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 25299fa..b826a8b 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -406,10 +406,22 @@ psyco_conn_tpc_recover(connectionObject *self, PyObject *dummy)
 static PyObject *
 psyco_conn_enter(connectionObject *self, PyObject *dummy)
 {
+    PyObject *rv = NULL;
+
     EXC_IF_CONN_CLOSED(self);
 
+    if (self->entered) {
+        PyErr_SetString(ProgrammingError,
+            "the connection cannot be re-entered recursively");
+        goto exit;
+    }
+
+    self->entered = 1;
     Py_INCREF(self);
-    return (PyObject *)self;
+    rv = (PyObject *)self;
+
+exit:
+    return rv;
 }
 
 
@@ -427,6 +439,9 @@ psyco_conn_exit(connectionObject *self, PyObject *args)
         goto exit;
     }
 
+    /* even if there will be an error, consider ourselves out */
+    self->entered = 0;
+
     if (type == Py_None) {
         if (!(tmp = PyObject_CallMethod((PyObject *)self, "commit", NULL))) {
             goto exit;
diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c
index 52ba78b..ff783ec 100644
--- a/psycopg/pqpath.c
+++ b/psycopg/pqpath.c
@@ -347,14 +347,19 @@ pq_begin_locked(connectionObject *conn, PyThreadState **tstate)
     char buf[256];  /* buf size must be same as bufsize */
     int result;
 
-    Dprintf("pq_begin_locked: pgconn = %p, autocommit = %d, status = %d",
+    Dprintf("pq_begin_locked: pgconn = %p, %d, status = %d",
             conn->pgconn, conn->autocommit, conn->status);
 
-    if (conn->autocommit || conn->status != CONN_STATUS_READY) {
+    if (conn->status != CONN_STATUS_READY) {
         Dprintf("pq_begin_locked: transaction in progress");
         return 0;
     }
 
+    if (conn->autocommit && !conn->entered) {
+        Dprintf("pq_begin_locked: autocommit and no with block");
+        return 0;
+    }
+
     if (conn->isolevel == ISOLATION_LEVEL_DEFAULT
             && conn->readonly == STATE_DEFAULT
             && conn->deferrable == STATE_DEFAULT) {
@@ -393,10 +398,10 @@ pq_commit(connectionObject *conn)
     Py_BEGIN_ALLOW_THREADS;
     pthread_mutex_lock(&conn->lock);
 
-    Dprintf("pq_commit: pgconn = %p, autocommit = %d, status = %d",
-            conn->pgconn, conn->autocommit, conn->status);
+    Dprintf("pq_commit: pgconn = %p, status = %d",
+            conn->pgconn, conn->status);
 
-    if (conn->autocommit || conn->status != CONN_STATUS_BEGIN) {
+    if (conn->status != CONN_STATUS_BEGIN) {
         Dprintf("pq_commit: no transaction to commit");
         retvalue = 0;
     }
@@ -427,10 +432,10 @@ pq_abort_locked(connectionObject *conn, PyThreadState **tstate)
 {
     int retvalue = -1;
 
-    Dprintf("pq_abort_locked: pgconn = %p, autocommit = %d, status = %d",
-            conn->pgconn, conn->autocommit, conn->status);
+    Dprintf("pq_abort_locked: pgconn = %p, status = %d",
+            conn->pgconn, conn->status);
 
-    if (conn->autocommit || conn->status != CONN_STATUS_BEGIN) {
+    if (conn->status != CONN_STATUS_BEGIN) {
         Dprintf("pq_abort_locked: no transaction to abort");
         return 0;
     }
@@ -488,12 +493,12 @@ pq_reset_locked(connectionObject *conn, PyThreadState **tstate)
 {
     int retvalue = -1;
 
-    Dprintf("pq_reset_locked: pgconn = %p, autocommit = %d, status = %d",
-            conn->pgconn, conn->autocommit, conn->status);
+    Dprintf("pq_reset_locked: pgconn = %p, status = %d",
+            conn->pgconn, conn->status);
 
     conn->mark += 1;
 
-    if (!conn->autocommit && conn->status == CONN_STATUS_BEGIN) {
+    if (conn->status == CONN_STATUS_BEGIN) {
         retvalue = pq_execute_command_locked(conn, "ABORT", tstate);
         if (retvalue != 0) return retvalue;
     }
diff --git a/tests/test_with.py b/tests/test_with.py
index 984602b..f52c747 100755
--- a/tests/test_with.py
+++ b/tests/test_with.py
@@ -155,6 +155,93 @@ class WithConnectionTestCase(WithTestCase):
         curs.execute("select * from test_with")
         self.assertEqual(curs.fetchall(), [])
 
+    def test_cant_reenter(self):
+        raised_ok = False
+        with self.conn:
+            try:
+                with self.conn:
+                    pass
+            except psycopg2.ProgrammingError:
+                raised_ok = True
+
+        self.assert_(raised_ok)
+
+        # Still good
+        with self.conn:
+            pass
+
+    def test_with_autocommit(self):
+        self.conn.autocommit = True
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+        with self.conn:
+            curs = self.conn.cursor()
+            curs.execute("insert into test_with values (1)")
+            self.assertEqual(
+                self.conn.info.transaction_status,
+                ext.TRANSACTION_STATUS_INTRANS,
+            )
+
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+        curs.execute("select count(*) from test_with")
+        self.assertEqual(curs.fetchone()[0], 1)
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+
+    def test_with_autocommit_pyerror(self):
+        self.conn.autocommit = True
+        raised_ok = False
+        try:
+            with self.conn:
+                curs = self.conn.cursor()
+                curs.execute("insert into test_with values (2)")
+                self.assertEqual(
+                    self.conn.info.transaction_status,
+                    ext.TRANSACTION_STATUS_INTRANS,
+                )
+                1 / 0
+        except ZeroDivisionError:
+            raised_ok = True
+
+        self.assert_(raised_ok)
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+        curs.execute("select count(*) from test_with")
+        self.assertEqual(curs.fetchone()[0], 0)
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+
+    def test_with_autocommit_pgerror(self):
+        self.conn.autocommit = True
+        raised_ok = False
+        try:
+            with self.conn:
+                curs = self.conn.cursor()
+                curs.execute("insert into test_with values (2)")
+                self.assertEqual(
+                    self.conn.info.transaction_status,
+                    ext.TRANSACTION_STATUS_INTRANS,
+                )
+                curs.execute("insert into test_with values ('x')")
+        except psycopg2.errors.InvalidTextRepresentation:
+            raised_ok = True
+
+        self.assert_(raised_ok)
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+        curs.execute("select count(*) from test_with")
+        self.assertEqual(curs.fetchone()[0], 0)
+        self.assertEqual(
+            self.conn.info.transaction_status, ext.TRANSACTION_STATUS_IDLE
+        )
+
 
 class WithCursorTestCase(WithTestCase):
     def test_with_ok(self):