Advanced rows count returning #221

For inlined INSERTs and from INSERT INTO ... SELECT

Author: xzkostyan

clickhouse_driver/client.py

@@ -255,8 +255,6 @@ def execute(self, query, params=None, with_column_types=False,
                          Defaults to ``False`` (row-like form).
 
         :return: * number of inserted rows for INSERT queries with data.
-                   Returning rows count from INSERT FROM SELECT is not
-                   supported.
                  * if `with_column_types=False`: `list` of `tuples` with
                    rows/columns.
                  * if `with_column_types=True`: `tuple` of 2 elements:
@@ -284,11 +282,23 @@ 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:
+                    # Dump check for INSERT.
+                    # Backwards compatibility for .execute(DDL) return [].
+                    # TODO: remove in 0.3.0. Return integer (zero) for DDL too.
+                    if query.lower().strip().startswith('insert'):
+                        rows = self.last_query.progress.rows
+                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,10 @@ def _process_response(self, response, executemany=False):
                 self._rowcount = len(rows)
         else:
             self._columns = self._types = []
+            # TODO: return 0 for DDL in 0.3.0
+            if not isinstance(rows, list):  # number of inserted rows
+                self._rowcount = rows
+                rows = []
 
         self._rows = 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_blocks.py

@@ -3,7 +3,7 @@
 
 from clickhouse_driver.errors import ServerException
 from tests.testcase import BaseTestCase, file_config
-from tests.util import capture_logging
+from tests.util import capture_logging, will_fail_in
 
 
 class BlocksTestCase(BaseTestCase):
@@ -80,6 +80,11 @@ def test_close_connection_on_keyboard_interrupt(self):
 
         self.assertFalse(self.client.connection.connected)
 
+    @will_fail_in(0, 3)
+    def test_ddl_return_value(self):
+        rv = self.client.execute('DROP TABLE test IF EXISTS')
+        self.assertEqual(rv, [])
+
 
 class ProgressTestCase(BaseTestCase):
     def test_select_with_progress(self):

tests/test_dbapi.py

@@ -8,6 +8,7 @@
     ProgrammingError, InterfaceError, OperationalError
 )
 from tests.testcase import BaseTestCase
+from tests.util import will_fail_in
 
 
 class DBAPITestCaseBase(BaseTestCase):
@@ -149,7 +150,15 @@ 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)
+            self.assertEqual(cursor.fetchall(), [])
+
+            cursor.execute(
+                'INSERT INTO test '
+                'SELECT number FROM system.numbers LIMIT 0'
+            )
+            self.assertEqual(cursor.rowcount, 0)
+            self.assertEqual(cursor.fetchall(), [])
 
     def test_description(self):
         with self.created_cursor() as cursor:
@@ -165,6 +174,7 @@ def test_pep249_sizes(self):
             cursor.setinputsizes(0)
             cursor.setoutputsize(0)
 
+    @will_fail_in(0, 3)
     def test_ddl(self):
         with self.created_cursor() as cursor:
             cursor.execute('DROP TABLE IF EXISTS test')

tests/test_insert.py

@@ -3,6 +3,7 @@
 from tests.testcase import BaseTestCase
 from clickhouse_driver import errors
 from clickhouse_driver.errors import ServerException
+from tests.util import require_server_version
 
 
 class InsertTestCase(BaseTestCase):
@@ -134,7 +135,21 @@ 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)
+
+    @require_server_version(19, 3, 3)
+    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'):

tests/util.py

@@ -2,6 +2,8 @@
 import logging
 from io import StringIO
 
+from clickhouse_driver import VERSION, __version__
+
 
 def skip_by_server_version(testcase, version_required):
     testcase.skipTest(
@@ -29,6 +31,21 @@ def wrapper(*args, **kwargs):
     return check
 
 
+def will_fail_in(*version):
+    def decorator(f):
+        @wraps(f)
+        def wrapper(*args, **kwargs):
+            self = args[0]
+
+            if VERSION >= version:
+                self.fail(
+                    'This test should not work in {}'.format(__version__)
+                )
+
+        return wrapper
+    return decorator
+
+
 class LoggingCapturer(object):
     def __init__(self, logger_name, level):
         self.old_stdout_handlers = []