[dbapi] Fixing header description

pydruid/db/api.py

@@ -206,21 +206,19 @@ def close(self):
     @check_closed
     def execute(self, operation, parameters=None):
         query = apply_parameters(operation, parameters or {})
-
         results = self._stream_query(query)
-        if self.header:
-            # The values are all null and thus it is impossible to imply types.
-            self.description = list(next(results)._asdict().items())
-            self._results = results
-        else:
-            # `_stream_query` returns a generator that produces the rows; we
-            # need to consume the first row so that `description` is properly
-            # set, so let's consume it and insert it back.
-            try:
-                first_row = next(results)
-                self._results = itertools.chain([first_row], results)
-            except StopIteration:
-                self._results = iter([])
+
+        # `_stream_query` returns a generator that produces the rows; we need to
+        # consume the first row so that `description` is properly set, so let's
+        # consume it and insert it back if it is not the header.
+        try:
+            first_row = next(results)
+            self._results = (
+                results if self.header else
+                itertools.chain([first_row], results)
+            )
+        except StopIteration:
+            self._results = iter([])
 
         return self
 
@@ -327,8 +325,11 @@ def _stream_query(self, query):
         Row = None
         for row in rows_from_chunks(chunks):
             # update description
-            if not self.header and self.description is None:
-                self.description = get_description_from_row(row)
+            if self.description is None:
+                self.description = (
+                    list(row.items()) if self.header else
+                    get_description_from_row(row)
+                )
 
             # return row in namedtuple
             if Row is None:

tests/db/test_cursor.py

@@ -103,6 +103,23 @@ def test_header_true(self, requests_post_mock):
         self.assertEquals(result, [Row(name='alice')])
         self.assertEquals(cursor.description, [('name', None)])
 
+    @patch('requests.post')
+    def test_names_with_underscores(self, requests_post_mock):
+        response = Response()
+        response.status_code = 200
+        response.raw = BytesIO(b'[{"_name": null}, {"_name": "alice"}]')
+        requests_post_mock.return_value = response
+        Row = namedtuple('Row', ['_name'], rename=True)
+
+        url = 'http://example.com/'
+        query = 'SELECT * FROM table'
+
+        cursor = Cursor(url, header=True)
+        cursor.execute(query)
+        result = cursor.fetchall()
+        self.assertEquals(result, [Row(_0='alice')])
+        self.assertEquals(cursor.description, [('_name', None)])
+
 
 if __name__ == '__main__':
     unittest.main()