feat: Allow setting of columnar and use_numpy for given_query

Author: paxcodes

clickhouse_backend/driver/connection.py

@@ -1,8 +1,11 @@
 import re
+import typing as T
+from contextlib import contextmanager
 
 from clickhouse_driver import connection
 from clickhouse_driver.dbapi import connection as dbapi_connection
 from clickhouse_driver.dbapi import cursor, errors
+from clickhouse_driver.result import IterQueryResult, ProgressQueryResult, QueryResult
 from django.conf import settings
 
 from .escape import escape_params
@@ -70,6 +73,11 @@ def send_query(self, query, query_id=None, params=None):
 
 
 class Cursor(cursor.Cursor):
+    
+    # Whether to return data in columnar format. For backwards-compatibility,
+    # let's default to None.
+    columnar = None
+
     def close(self):
         """Push client back to connection pool"""
         if self.closed:
@@ -81,12 +89,64 @@ def close(self):
     def closed(self):
         return self._state == self._states.CURSOR_CLOSED
 
+    @property
+    def use_numpy(self):
+        return self._client.client_settings["use_numpy"]
+
+    @use_numpy.setter
+    def use_numpy(self, value):
+        self._client.client_settings["use_numpy"] = value
+        if value:
+            try:
+                from clickhouse_driver.numpy.result import (
+                    NumpyIterQueryResult,
+                    NumpyProgressQueryResult,
+                    NumpyQueryResult,
+                )
+
+                self._client.query_result_cls = NumpyQueryResult
+                self._client.iter_query_result_cls = NumpyIterQueryResult
+                self._client.progress_query_result_cls = NumpyProgressQueryResult
+            except ImportError as e:
+                raise RuntimeError("Extras for NumPy must be installed") from e
+        else:
+            self._client.query_result_cls = QueryResult
+            self._client.iter_query_result_cls = IterQueryResult
+            self._client.progress_query_result_cls = ProgressQueryResult
+
+    @contextmanager
+    def set_query_args(
+        self, columnar: T.Optional[bool] = None, use_numpy: T.Optional[bool] = None
+    ):
+        original_use_numpy = self.use_numpy
+        if use_numpy is not None:
+            self.use_numpy = use_numpy
+
+        original_columnar = self.columnar
+        if columnar is not None:
+            self.columnar = columnar
+
+        yield self
+
+        self.use_numpy = original_use_numpy
+        self.columnar = original_columnar
+
     def __del__(self):
         # If someone forgets calling close method,
         # then release connection when gc happens.
         if not self.closed:
             self.close()
 
+    def _prepare(self):
+        """Override clickhouse_driver.Cursor._prepare() to add columnar kwargs.
+
+        See https://github.com/jayvynl/django-clickhouse-backend/issues/119
+        """
+        execute, execute_kwargs = super()._prepare()
+        if self.columnar is not None:
+            execute_kwargs["columnar"] = self.columnar
+        return execute, execute_kwargs
+
     def execute(self, operation, parameters=None):
         """fix https://github.com/jayvynl/django-clickhouse-backend/issues/9"""
         if getattr(