Advanced rows count returning #221

For inlined INSERTs and from INSERT INTO ... SELECT

clickhouse_driver/client.py

@@ -284,11 +284,20 @@ def execute(self, query, params=None, with_column_types=False,
                 )
             else:
                 rv = self.process_ordinary_query(
-                    query, params=params, with_column_types=with_column_types,
+                    query, params=params, with_column_types=True,
                     external_tables=external_tables,
                     query_id=query_id, types_check=types_check,
                     columnar=columnar
                 )
+                rows, columns_with_types = rv
+                # no columns in case of DDL or INSERT ... SELECT
+                if not columns_with_types:
+                    progress = self.last_query.progress
+                    rows = progress.rows if progress._is_incremented else []
+                rv = (rows, columns_with_types)
+                if not with_column_types:
+                    rv = rv[0]
+
             self.last_query.store_elapsed(time() - start_time)
             return rv
 

clickhouse_driver/dbapi/cursor.py

@@ -332,6 +332,9 @@ def _process_response(self, response, executemany=False):
                 self._rowcount = len(rows)
         else:
             self._columns = self._types = []
+            if not (isinstance(rows, list) and not len(rows)):
+                # number of inserted rows
+                self._rowcount = rows
 
         self._rows = rows
 

clickhouse_driver/progress.py

@@ -9,6 +9,7 @@ def __init__(self):
         self.total_rows = 0
         self.written_rows = 0
         self.written_bytes = 0
+        self._is_incremented = False
 
         super(Progress, self).__init__()
 
@@ -25,6 +26,7 @@ def read(self, server_revision, fin):
             self.written_bytes = read_varint(fin)
 
     def increment(self, another_progress):
+        self._is_incremented = True
         self.rows += another_progress.rows
         self.bytes += another_progress.bytes
         self.total_rows += another_progress.total_rows

docs/features.rst

@@ -324,7 +324,7 @@ Parameters are expected in Python extended format codes, e.g.
         ...     {'limit': 3}
         ... )
         >>> cursor.rowcount
-        0
+        5
         >>> cursor.execute('SELECT sum(x) FROM test')
         >>> cursor.fetchall()
         [(303,)]

docs/quickstart.rst

@@ -193,7 +193,7 @@ Of course for ``INSERT ... SELECT`` queries data is not needed:
         ...     'SELECT * FROM system.numbers LIMIT %(limit)s',
         ...     {'limit': 5}
         ... )
-        []
+        5
 
 ClickHouse will execute this query like a usual ``SELECT`` query.
 

tests/test_dbapi.py

@@ -149,7 +149,12 @@ def test_rowcount_insert_from_select(self):
                 'INSERT INTO test '
                 'SELECT number FROM system.numbers LIMIT 4'
             )
-            self.assertEqual(cursor.rowcount, -1)
+            self.assertEqual(cursor.rowcount, 4)
+            cursor.execute(
+                'INSERT INTO test '
+                'SELECT number FROM system.numbers LIMIT 0'
+            )
+            self.assertEqual(cursor.rowcount, 0)
 
     def test_description(self):
         with self.created_cursor() as cursor:

tests/test_insert.py

@@ -134,7 +134,20 @@ def test_insert_from_select(self):
                 'INSERT INTO test (a) '
                 'SELECT number FROM system.numbers LIMIT 5'
             )
-            self.assertEqual(inserted, [])
+            self.assertEqual(inserted, 5)
+
+            inserted = self.client.execute(
+                'INSERT INTO test (a) '
+                'SELECT number FROM system.numbers LIMIT 0'
+            )
+            self.assertEqual(inserted, 0)
+
+    def test_insert_inline(self):
+        with self.create_table('a UInt64'):
+            inserted = self.client.execute(
+                'INSERT INTO test (a) VALUES (1), (2), (3)'
+            )
+            self.assertEqual(inserted, 3)
 
     def test_insert_return(self):
         with self.create_table('a Int8'):