Commit576eafc

authored

Fix potential resource leak inCloudFetchQueue (#624)

* add close() for Queue, add ResultSet invocationSigned-off-by: varun-edachali-dbx <varun.edachali@databricks.com>* move Queue closure to finally: block to ensure client-side cleanup regardless of server side stateSigned-off-by: varun-edachali-dbx <varun.edachali@databricks.com>* add unit test assertions to ensure Queue closureSigned-off-by: varun-edachali-dbx <varun.edachali@databricks.com>* move results closure to try blockSigned-off-by: varun-edachali-dbx <varun.edachali@databricks.com>---------Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

1 parent9c34acd commit576eafcCopy full SHA for 576eafc

File tree

3 files changed

+20

-0

lines changed

src/databricks/sql
- client.py
- utils.py
tests/unit
- test_client.py

3 files changed

+20

-0

lines changed

`‎src/databricks/sql/client.py‎`

Lines changed: 1 addition & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -1596,6 +1596,7 @@ def close(self) -> None:`
`1596`	`1596`	`been closed on the server for some other reason, issue a request to the server to close it.`
`1597`	`1597`	`"""`
`1598`	`1598`	`try:`
	`1599`	`+self.results.close()`
`1599`	`1600`	`if (`
`1600`	`1601`	`self.op_state!=self.thrift_backend.CLOSED_OP_STATE`
`1601`	`1602`	`andnotself.has_been_closed_server_side`

`‎src/databricks/sql/utils.py‎`

Lines changed: 13 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,10 @@ def next_n_rows(self, num_rows: int):`
`46`	`46`	`defremaining_rows(self):`
`47`	`47`	`pass`
`48`	`48`
	`49`	`+@abstractmethod`
	`50`	`+defclose(self):`
	`51`	`+pass`
	`52`	`+`
`49`	`53`
`50`	`54`	`classResultSetQueueFactory(ABC):`
`51`	`55`	`@staticmethod`
`@@ -157,6 +161,9 @@ def remaining_rows(self):`
`157`	`161`	`self.cur_row_index+=slice.num_rows`
`158`	`162`	`returnslice`
`159`	`163`
	`164`	`+defclose(self):`
	`165`	`+return`
	`166`	`+`
`160`	`167`
`161`	`168`	`classArrowQueue(ResultSetQueue):`
`162`	`169`	`def__init__(`
`@@ -192,6 +199,9 @@ def remaining_rows(self) -> "pyarrow.Table":`
`192`	`199`	`self.cur_row_index+=slice.num_rows`
`193`	`200`	`returnslice`
`194`	`201`
	`202`	`+defclose(self):`
	`203`	`+return`
	`204`	`+`
`195`	`205`
`196`	`206`	`classCloudFetchQueue(ResultSetQueue):`
`197`	`207`	`def__init__(`
`@@ -341,6 +351,9 @@ def _create_empty_table(self) -> "pyarrow.Table":`
`341`	`351`	`# Create a 0-row table with just the schema bytes`
`342`	`352`	`returncreate_arrow_table_from_arrow_file(self.schema_bytes,self.description)`
`343`	`353`
	`354`	`+defclose(self):`
	`355`	`+self.download_manager._shutdown_manager()`
	`356`	`+`
`344`	`357`
`345`	`358`	`ExecuteResponse=namedtuple(`
`346`	`359`	`"ExecuteResponse",`

`‎tests/unit/test_client.py‎`

Lines changed: 6 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -267,33 +267,39 @@ def test_arraysize_buffer_size_passthrough(`
`267`	`267`	`deftest_closing_result_set_with_closed_connection_soft_closes_commands(self):`
`268`	`268`	`mock_connection=Mock()`
`269`	`269`	`mock_backend=Mock()`
	`270`	`+mock_results=Mock()`
`270`	`271`	`result_set=client.ResultSet(`
`271`	`272`	`connection=mock_connection,`
`272`	`273`	`thrift_backend=mock_backend,`
`273`	`274`	`execute_response=Mock(),`
`274`	`275`	`)`
	`276`	`+result_set.results=mock_results`
`275`	`277`	`mock_connection.open=False`
`276`	`278`
`277`	`279`	`result_set.close()`
`278`	`280`
`279`	`281`	`self.assertFalse(mock_backend.close_command.called)`
`280`	`282`	`self.assertTrue(result_set.has_been_closed_server_side)`
	`283`	`+mock_results.close.assert_called_once()`
`281`	`284`
`282`	`285`	`deftest_closing_result_set_hard_closes_commands(self):`
`283`	`286`	`mock_results_response=Mock()`
`284`	`287`	`mock_results_response.has_been_closed_server_side=False`
`285`	`288`	`mock_connection=Mock()`
`286`	`289`	`mock_thrift_backend=Mock()`
	`290`	`+mock_results=Mock()`
`287`	`291`	`mock_connection.open=True`
`288`	`292`	`result_set=client.ResultSet(`
`289`	`293`	`mock_connection,mock_results_response,mock_thrift_backend`
`290`	`294`	`)`
	`295`	`+result_set.results=mock_results`
`291`	`296`
`292`	`297`	`result_set.close()`
`293`	`298`
`294`	`299`	`mock_thrift_backend.close_command.assert_called_once_with(`
`295`	`300`	`mock_results_response.command_handle`
`296`	`301`	`)`
	`302`	`+mock_results.close.assert_called_once()`
`297`	`303`
`298`	`304`	`@patch("%s.client.ResultSet"%PACKAGE_NAME)`
`299`	`305`	`deftest_executing_multiple_commands_uses_the_most_recent_command(`

0 commit comments

Comments

(0)

Movatterモバイル変換

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Commit576eafc

File tree

3 files changed

3 files changed

`‎src/databricks/sql/client.py‎`

`‎src/databricks/sql/utils.py‎`

`‎tests/unit/test_client.py‎`

0 commit comments