diff --git a/clickhouse_driver/dbapi/cursor.py b/clickhouse_driver/dbapi/cursor.py index abfad408..5e5574d9 100644 --- a/clickhouse_driver/dbapi/cursor.py +++ b/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. diff --git a/tests/test_dbapi.py b/tests/test_dbapi.py index 49f0163b..1b7cf681 100644 --- a/tests/test_dbapi.py +++ b/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, )]