diff --git a/pydruid/db/api.py b/pydruid/db/api.py index 4d5902c0..40b4a0e4 100644 --- a/pydruid/db/api.py +++ b/pydruid/db/api.py @@ -186,8 +186,11 @@ def execute(self, operation, parameters=None): # to consume the first row so that `description` is properly set, so # let's consume it and insert it back. results = self._stream_query(query) - first_row = next(results) - self._results = itertools.chain([first_row], results) + try: + first_row = next(results) + self._results = itertools.chain([first_row], results) + except StopIteration: + self._results = iter([]) return self diff --git a/tests/db/test_cursor.py b/tests/db/test_cursor.py new file mode 100644 index 00000000..f1d6599f --- /dev/null +++ b/tests/db/test_cursor.py @@ -0,0 +1,48 @@ +# -*- coding: utf-8 -*- + +from collections import namedtuple +from mock import patch +import unittest + +from requests.models import Response +from six import BytesIO + +from pydruid.db.api import Cursor + + +class CursorTestSuite(unittest.TestCase): + + @patch('requests.post') + def test_execute(self, requests_post_mock): + response = Response() + response.status_code = 200 + response.raw = BytesIO(b'[{"name": "alice"}, {"name": "bob"}, {"name": "charlie"}]') + requests_post_mock.return_value = response + Row = namedtuple('Row', ['name']) + + cursor = Cursor('http://example.com/') + cursor.execute('SELECT * FROM table') + result = cursor.fetchall() + expected = [ + Row(name='alice'), + Row(name='bob'), + Row(name='charlie'), + ] + self.assertEquals(result, expected) + + @patch('requests.post') + def test_execute_empty_result(self, requests_post_mock): + response = Response() + response.status_code = 200 + response.raw = BytesIO(b'[]') + requests_post_mock.return_value = response + + cursor = Cursor('http://example.com/') + cursor.execute('SELECT * FROM table') + result = cursor.fetchall() + expected = [] + self.assertEquals(result, expected) + + +if __name__ == '__main__': + unittest.main()