Properly cleanup memory of broken connections

Fixed ticket #148.

NEWS

@@ -1,3 +1,9 @@
+What's new in psycopg 2.4.7
+---------------------------
+
+  - Properly cleanup memory of broken connections (ticket #142).
+
+
 What's new in psycopg 2.4.6
 ---------------------------
 

psycopg/connection_int.c

@@ -918,7 +918,8 @@ conn_poll(connectionObject *self)
 void
 conn_close(connectionObject *self)
 {
-    if (self->closed) {
+    /* a connection with closed == 2 still requires cleanup */
+    if (self->closed == 1) {
         return;
     }
 
@@ -936,7 +937,7 @@ conn_close(connectionObject *self)
 
 void conn_close_locked(connectionObject *self)
 {
-    if (self->closed) {
+    if (self->closed == 1) {
         return;
     }
 
@@ -957,6 +958,8 @@ void conn_close_locked(connectionObject *self)
         PQfinish(self->pgconn);
         self->pgconn = NULL;
         Dprintf("conn_close: PQfinish called");
+    }
+    if (self->cancel) {
         PQfreeCancel(self->cancel);
         self->cancel = NULL;
     }

psycopg/connection_type.c

@@ -1068,7 +1068,7 @@ connection_dealloc(PyObject* obj)
 
     PyObject_GC_UnTrack(self);
 
-    if (self->closed == 0) conn_close(self);
+    conn_close(self);
 
     conn_notice_clean(self);
 

tests/test_connection.py

@@ -26,10 +26,11 @@ import os
 import time
 import threading
 from testutils import unittest, decorate_all_tests
-from testutils import skip_before_postgres, skip_after_postgres
+from testutils import skip_before_postgres, skip_after_postgres, skip_if_no_superuser
 from operator import attrgetter
 
 import psycopg2
+import psycopg2.errorcodes
 import psycopg2.extensions
 from testconfig import dsn, dbname
 
@@ -66,6 +67,22 @@ class ConnectionTests(unittest.TestCase):
         conn.close()
         self.assertEqual(curs.closed, True)
 
+    @skip_before_postgres(8, 4)
+    @skip_if_no_superuser
+    def test_cleanup_on_badconn_close(self):
+        # ticket #148
+        conn = self.conn
+        cur = conn.cursor()
+        try:
+            cur.execute("select pg_terminate_backend(pg_backend_pid())")
+        except psycopg2.OperationalError, e:
+            if e.pgcode != psycopg2.errorcodes.ADMIN_SHUTDOWN:
+                raise
+
+        self.assertEqual(conn.closed, 2)
+        conn.close()
+        self.assertEqual(conn.closed, 1)
+
     def test_reset(self):
         conn = self.conn
         # switch isolation level, then reset

tests/testutils.py

@@ -207,6 +207,21 @@ def skip_from_python(*ver):
         return skip_from_python__
     return skip_from_python_
 
+def skip_if_no_superuser(f):
+    """Skip a test if the database user running the test is not a superuser"""
+    def skip_if_no_superuser_(self):
+        from psycopg2 import ProgrammingError
+        try:
+            return f(self)
+        except ProgrammingError, e:
+            import psycopg2.errorcodes
+            if e.pgcode == psycopg2.errorcodes.INSUFFICIENT_PRIVILEGE:
+                self.skipTest("skipped because not superuser")
+            else:
+                raise
+
+    return skip_if_no_superuser_
+
 
 def script_to_py3(script):
     """Convert a script to Python3 syntax if required."""