Iterator + context managerAPI

This small PR adds APIs to the Python client for iteration and context managers.

New unit tests

Author: Niall Egan <niall.egan@databricks.com>

Author: Niall Egan

cmdexec/clients/python/command_exec_client.py

@@ -42,6 +42,12 @@ def __init__(self, **kwargs):
         logger.info("Successfully opened session " + str(self.session_id.hex()))
         self._cursors = []
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def cursor(self, buffer_size_rows=DEFAULT_BUFFER_SIZE_ROWS):
         if not self.open:
             raise Error("Cannot create cursor from closed connection")
@@ -69,6 +75,19 @@ def __init__(self, connection, buffer_size_rows=DEFAULT_BUFFER_SIZE_ROWS):
         # Note that Cursor closed => active result set closed, but not vice versa
         self.open = True
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def __iter__(self):
+        if self.active_result_set:
+            for row in self.active_result_set:
+                yield row
+        else:
+            raise Error("There is no active result set")
+
     def _response_to_result_set(self, execute_command_response, status):
         command_id = execute_command_response.command_id
         arrow_results = execute_command_response.results.arrow_ipc_stream
@@ -199,6 +218,14 @@ def __init__(self,
             # In this case, there are results waiting on the server so we fetch now for simplicity
             self._fill_results_buffer()
 
+    def __iter__(self):
+        while True:
+            row = self.fetchone()
+            if row:
+                yield row
+            else:
+                break
+
     def _fetch_and_deserialize_results(self):
         fetch_results_request = command_pb2.FetchCommandResultsRequest(
             id=self.command_id,

cmdexec/clients/python/tests/tests.py

@@ -176,6 +176,25 @@ def test_negative_fetch_throws_exception(self, pyarrow_ipc_open_stream_mock):
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)
 
+    def test_context_manager_closes_cursor(self):
+        mock_close = Mock()
+        with command_exec_client.Cursor(Mock()) as cursor:
+            cursor.close = mock_close
+        mock_close.assert_called_once_with()
+
+    @patch("cmdexec.clients.python.command_exec_client.CmdExecBaseHttpClient")
+    def test_context_manager_closes_connection(self, mock_client_class):
+        instance = mock_client_class.return_value
+        mock_response = Mock()
+        mock_response.id = b'\x22'
+        instance.make_request.return_value = mock_response
+        good_connection_args = {"HOST": 1, "PORT": 1}
+        mock_close = Mock()
+
+        with command_exec_client.connect(**good_connection_args) as connection:
+            connection.close = mock_close
+        mock_close.assert_called_once_with()
+
 
 if __name__ == '__main__':
     suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])