acquire lock before notif + formatting (black)

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/queue.py

@@ -278,7 +278,8 @@ def _worker_loop(self):
             if not links_downloaded:
                 self._shutdown_event.set()
         logger.debug("LinkFetcher[%s]: worker thread exiting", self._statement_id)
-        self._link_data_update.notify_all()
+        with self._link_data_update:
+            self._link_data_update.notify_all()
 
     def start(self):
         """Spawn the worker thread."""

tests/unit/test_client.py

@@ -224,7 +224,10 @@ def test_closing_result_set_hard_closes_commands(self):
 
         mock_thrift_backend.fetch_results.return_value = (Mock(), False, 0)
         result_set = ThriftResultSet(
-            mock_connection, mock_results_response, mock_thrift_backend, session_id_hex=Mock()
+            mock_connection,
+            mock_results_response,
+            mock_thrift_backend,
+            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -272,7 +275,9 @@ def test_negative_fetch_throws_exception(self):
         mock_backend = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-        result_set = ThriftResultSet(Mock(), Mock(), mock_backend, session_id_hex=Mock())
+        result_set = ThriftResultSet(
+            Mock(), Mock(), mock_backend, session_id_hex=Mock()
+        )
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)

tests/unit/test_downloader.py

@@ -27,7 +27,12 @@ def test_run_link_expired(self, mock_time):
         # Already expired
         result_link.expiryTime = 999
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
 
         with self.assertRaises(Error) as context:
@@ -43,7 +48,12 @@ def test_run_link_past_expiry_buffer(self, mock_time):
         # Within the expiry buffer time
         result_link.expiryTime = 1004
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
 
         with self.assertRaises(Error) as context:
@@ -63,7 +73,12 @@ def test_run_get_response_not_ok(self, mock_time, mock_session):
         result_link = Mock(expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
         with self.assertRaises(requests.exceptions.HTTPError) as context:
             d.run()
@@ -82,7 +97,12 @@ def test_run_uncompressed_successful(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
         file = d.run()
 
@@ -105,7 +125,12 @@ def test_run_compressed_successful(self, mock_time, mock_session):
         result_link = Mock(bytesNum=100, expiryTime=1001)
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
         file = d.run()
 
@@ -121,7 +146,12 @@ def test_download_connection_error(self, mock_time, mock_session):
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
         with self.assertRaises(ConnectionError):
             d.run()
@@ -136,7 +166,12 @@ def test_download_timeout(self, mock_time, mock_session):
         mock_session.return_value.get.return_value.content = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'
 
         d = downloader.ResultSetDownloadHandler(
-            settings, result_link, ssl_options=SSLOptions(), chunk_id=0, session_id_hex=Mock(), statement_id=Mock()
+            settings,
+            result_link,
+            ssl_options=SSLOptions(),
+            chunk_id=0,
+            session_id_hex=Mock(),
+            statement_id=Mock(),
         )
         with self.assertRaises(TimeoutError):
             d.run()

tests/unit/test_thrift_backend.py

@@ -731,7 +731,9 @@ def test_get_status_uses_display_message_if_available(self, tcli_service_class):
             ssl_options=SSLOptions(),
         )
         with self.assertRaises(DatabaseError) as cm:
-            thrift_backend.execute_command(Mock(), Mock(), 100, 100, Mock(), Mock(), Mock())
+            thrift_backend.execute_command(
+                Mock(), Mock(), 100, 100, Mock(), Mock(), Mock()
+            )
 
         self.assertEqual(display_message, str(cm.exception))
         self.assertIn(diagnostic_info, str(cm.exception.message_with_context()))
@@ -772,7 +774,9 @@ def test_direct_results_uses_display_message_if_available(self, tcli_service_cla
             ssl_options=SSLOptions(),
         )
         with self.assertRaises(DatabaseError) as cm:
-            thrift_backend.execute_command(Mock(), Mock(), 100, 100, Mock(), Mock(), Mock())
+            thrift_backend.execute_command(
+                Mock(), Mock(), 100, 100, Mock(), Mock(), Mock()
+            )
 
         self.assertEqual(display_message, str(cm.exception))
         self.assertIn(diagnostic_info, str(cm.exception.message_with_context()))
@@ -1450,7 +1454,9 @@ def test_non_arrow_non_column_based_set_triggers_exception(
         thrift_backend = self._make_fake_thrift_backend()
 
         with self.assertRaises(OperationalError) as cm:
-            thrift_backend.execute_command("foo", Mock(), 100, 100, Mock(), Mock(), Mock())
+            thrift_backend.execute_command(
+                "foo", Mock(), 100, 100, Mock(), Mock(), Mock()
+            )
         self.assertIn(
             "Expected results to be in Arrow or column based format", str(cm.exception)
         )
@@ -2279,7 +2285,9 @@ def test_execute_command_sets_complex_type_fields_correctly(
                 ssl_options=SSLOptions(),
                 **complex_arg_types,
             )
-            thrift_backend.execute_command(Mock(), Mock(), 100, 100, Mock(), Mock(), Mock())
+            thrift_backend.execute_command(
+                Mock(), Mock(), 100, 100, Mock(), Mock(), Mock()
+            )
             t_execute_statement_req = tcli_service_instance.ExecuteStatement.call_args[
                 0
             ][0]