Add cursor attribute columns_with_types to DB API

clickhouse_driver/dbapi/cursor.py

@@ -29,6 +29,10 @@ def __init__(self, client):
 
         self.arraysize = 1
 
+        # Begin non-PEP attributes
+        self._columns_with_types = None
+        # End non-PEP attributes
+
         super(Cursor, self).__init__()
 
     def __repr__(self):
@@ -206,6 +210,14 @@ def setoutputsize(self, size, column=None):
         pass
 
     # Begin non-PEP methods
+    @property
+    def columns_with_types(self):
+        """
+        :return: list of column names with corresponding types of the last
+                 .execute*(). E.g. [('x', 'UInt64')].
+        """
+        return self._columns_with_types
+
     def set_stream_results(self, stream_results, max_row_buffer):
         """
         Toggles results streaming from server. Driver will consume
@@ -302,6 +314,8 @@ def _process_response(self, response, executemany=False):
         else:
             rows, columns_with_types = response
 
+        self._columns_with_types = columns_with_types
+
         # Only SELECT queries have columns_with_types.
         # DDL and INSERT INTO ... SELECT queries have empty columns header.
         # We need to obtain rows count only during non-streaming SELECTs.

tests/test_dbapi.py

@@ -214,6 +214,32 @@ def test_connection_repr(self):
 
 
 class ExtrasTestCase(DBAPITestCaseBase):
+    def test_columns_with_types_select(self):
+        with self.created_cursor() as cursor:
+            self.assertIsNone(cursor.columns_with_types)
+            cursor.execute(
+                'SELECT CAST(number AS UInt64) AS x '
+                'FROM system.numbers LIMIT 4'
+            )
+            cursor.fetchall()
+            self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
+
+    def test_columns_with_types_insert(self):
+        with self.created_cursor() as cursor, self.create_table('a UInt8'):
+            cursor.executemany('INSERT INTO test (a) VALUES', [(123, )])
+            self.assertIsNone(cursor.columns_with_types)
+
+    def test_columns_with_types_streaming(self):
+        with self.created_cursor() as cursor:
+            cursor.set_stream_results(True, 2)
+            cursor.execute(
+                'SELECT CAST(number AS UInt64) AS x '
+                'FROM system.numbers LIMIT 4'
+            )
+            self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
+            list(cursor)
+            self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
+
     def test_set_external_tables(self):
         with self.created_cursor() as cursor:
             data = [(0, ), (1, ), (2, )]