fix: don't try to close closed cursors

googleapis · Jan 28, 2021 · df9776b · df9776b
1 parent d5735ea
commit df9776b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/google/cloud/bigquery/dbapi/connection.py b/google/cloud/bigquery/dbapi/connection.py
@@ -76,7 +76,8 @@ def close(self):
             self._bqstorage_client._transport.grpc_channel.close()
 
         for cursor_ in self._cursors_created:
-            cursor_.close()
+            if not cursor_._closed:
+                cursor_.close()
 
     def commit(self):
         """No-op, but for consistency raise an error if connection is closed."""

diff --git a/tests/unit/test_dbapi_connection.py b/tests/unit/test_dbapi_connection.py
@@ -67,7 +67,8 @@ def test_ctor_w_bqstorage_client(self):
         mock_client = self._mock_client()
         mock_bqstorage_client = self._mock_bqstorage_client()
         connection = self._make_one(
-            client=mock_client, bqstorage_client=mock_bqstorage_client,
+            client=mock_client,
+            bqstorage_client=mock_bqstorage_client,
         )
         self.assertIsInstance(connection, Connection)
         self.assertIs(connection._client, mock_client)
@@ -109,7 +110,8 @@ def test_connect_w_both_clients(self):
         mock_client = self._mock_client()
         mock_bqstorage_client = self._mock_bqstorage_client()
         connection = connect(
-            client=mock_client, bqstorage_client=mock_bqstorage_client,
+            client=mock_client,
+            bqstorage_client=mock_bqstorage_client,
         )
         self.assertIsInstance(connection, Connection)
         self.assertIs(connection._client, mock_client)
@@ -140,7 +142,9 @@ def test_close_closes_all_created_bigquery_clients(self):
             return_value=client,
         )
         bqstorage_client_patcher = mock.patch.object(
-            client, "_create_bqstorage_client", return_value=bqstorage_client,
+            client,
+            "_create_bqstorage_client",
+            return_value=bqstorage_client,
         )
 
         with client_patcher, bqstorage_client_patcher:
@@ -156,7 +160,7 @@ def test_close_closes_all_created_bigquery_clients(self):
     )
     def test_close_does_not_close_bigquery_clients_passed_to_it(self):
         client = self._mock_client()
-        bqstorage_client = self._mock_bqstorage_client()
+        bqstorage_client = sesf._mock_bqstorage_client()
         connection = self._make_one(client=client, bqstorage_client=bqstorage_client)
 
         connection.close()
@@ -176,6 +180,19 @@ def test_close_closes_all_created_cursors(self):
         self.assertTrue(cursor_1._closed)
         self.assertTrue(cursor_2._closed)
 
+    def test_close_closes_only_open_created_cursors(self):
+        connection = self._make_one(client=self._mock_client())
+        cursor_1 = connection.cursor()
+        cursor_2 = connection.cursor()
+        self.assertFalse(cursor_1._closed)
+        self.assertFalse(cursor_2._closed)
+
+        cursor_1.close()
+        connection.close()
+
+        self.assertTrue(cursor_1._closed)
+        self.assertTrue(cursor_2._closed)
+
     def test_does_not_keep_cursor_instances_alive(self):
         from google.cloud.bigquery.dbapi import Cursor