Stop reusing broken connection

Author: caitriona-cloudsmith

clickhouse_backend/driver/pool.py

@@ -105,6 +105,11 @@ def push(
                     raise InterfaceError("trying to put unkeyed client")
             if len(self._pool) < self.connections_min and not close:
                 # TODO: verify connection still valid
+
+                # If the connection is currently executing a query, it shouldn't be reused.
+                # Explicitly disconnect it instead.
+                if client.connection.is_query_executing:
+                    client.disconnect()
                 if client.connection.connected:
                     self._pool.append(client)
             else:

tests/backends/clickhouse/test_driver.py

@@ -1,9 +1,4 @@
-from clickhouse_driver.dbapi.errors import (
-    OperationalError as ch_driver_OperationalError,
-)
-from clickhouse_driver.errors import PartiallyConsumedQueryError
 from django.db import connection
-from django.db.utils import OperationalError as django_OperationalError
 from django.test import TestCase
 
 from clickhouse_backend.driver import connect
@@ -34,41 +29,19 @@ def setUpTestData(cls):
             ]
         )
 
-    def test_connection_unusable_when_iteration_interrupted(self):
-        """
-        This test demonstrates that if a queryset is iterated over and the iteration
-        is interrupted (e.g. via a break statement), the connection used for that
-        iteration is not cleaned up and is left in a broken state. Any subsequent
-        queries using that connection will fail.
-        """
+    def test_connection_not_reused_when_iteration_interrupted(self):
         pool = connection.connection.pool
-        connection_count_before = len(pool._pool)
 
+        connection_count_before = len(pool._pool)
         assert connection_count_before == 1
 
-        # Asserts most recent exception is Django OperationalError
-        with self.assertRaises(django_OperationalError) as ex_context:
-            # Get queryset
-            authors = models.Author.objects.all()
-            # Access iterator, but break after first item
-            for author in authors.iterator(1):
-                author = author.name
-                break
-
-            # Assert connection pool size is unchanged despite broken connection
-            connection_count_after_iterator = len(pool._pool)
-            assert connection_count_after_iterator == 1
+        authors = models.Author.objects.all()
+        for author in authors.iterator(1):
+            author = author.name
+            break
 
-            # Try to access queryset again, which won't work via same connection
-            author = authors.get(id=self.a1.id)
+        connection_count_after_iterator = len(pool._pool)
+        # Connection was closed and not returned to pool
+        assert connection_count_after_iterator == 0
 
-        # Caused by ch driver driver Operational error
-        self.assertIsInstance(
-            ex_context.exception.__cause__, ch_driver_OperationalError
-        )
-
-        # ...The context of which is a PartiallyConsumedQueryError
-        # https://github.com/mymarilyn/clickhouse-driver/blob/master/clickhouse_driver/connection.py#L801
-        self.assertIsInstance(
-            ex_context.exception.__cause__.__context__, PartiallyConsumedQueryError
-        )
+        author = authors.get(id=self.a1.id)