Intermediate changes

8 files changed, 75 insertions, 35 deletions

diff --git a/contrib/python/clickhouse-connect/.dist-info/METADATA b/contrib/python/clickhouse-connect/.dist-info/METADATA
index 914107f75a4..0d75e95b0e4 100644
--- a/contrib/python/clickhouse-connect/.dist-info/METADATA
+++ b/contrib/python/clickhouse-connect/.dist-info/METADATA
@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: clickhouse-connect
-Version: 0.7.6
+Version: 0.7.7
 Summary: ClickHouse Database Core Driver for Python, Pandas, and Superset
 Home-page: https://github.com/ClickHouse/clickhouse-connect
 Author: ClickHouse Inc.
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/__version__.py b/contrib/python/clickhouse-connect/clickhouse_connect/__version__.py
index 4a603d90096..9e59d8f8e18 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/__version__.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/__version__.py
@@ -1 +1 @@
-version = '0.7.6'
+version = '0.7.7'
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py
index 14d60351f42..0c0d25e6b0a 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/cc_sqlalchemy/datatypes/base.py
@@ -5,7 +5,7 @@ from sqlalchemy.exc import CompileError
 
 from clickhouse_connect.datatypes.base import ClickHouseType, TypeDef, EMPTY_TYPE_DEF
 from clickhouse_connect.datatypes.registry import parse_name, type_map
-from clickhouse_connect.driver.query import format_query_value
+from clickhouse_connect.driver.query import str_query_value
 
 logger = logging.getLogger(__name__)
 
@@ -96,12 +96,12 @@ class ChSqlaType:
         method and should be able to ignore literal_processor definitions in the dialect, which are verbose and
         confusing.
         """
-        return format_query_value
+        return str_query_value
 
     def _compiler_dispatch(self, _visitor, **_):
         """
         Override for the SqlAlchemy TypeEngine _compiler_dispatch method to sidestep unnecessary layers and complexity
-        when generating the type name.  The underlying ClickHouseType generates the correct name
+        when generating the type name.  The underlying ClickHouseType generates the correct name for the type
         :return: Name generated by the underlying driver.
         """
         return self.name
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py
index cf16ec24ece..3cce716372d 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/client.py
@@ -256,8 +256,7 @@ class Client(ABC):
                   settings: Optional[Dict[str, Any]] = None,
                   fmt: str = None,
                   use_database: bool = True,
-                  external_data: Optional[ExternalData] = None,
-                  stream: bool = False) -> Union[bytes, io.IOBase]:
+                  external_data: Optional[ExternalData] = None) -> bytes:
         """
         Query method that simply returns the raw ClickHouse format bytes
         :param query: Query statement/format string
@@ -270,6 +269,25 @@ class Client(ABC):
         :return: bytes representing raw ClickHouse return value based on format
         """
 
+    @abstractmethod
+    def raw_stream(self, query: str,
+                   parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
+                   settings: Optional[Dict[str, Any]] = None,
+                   fmt: str = None,
+                   use_database: bool = True,
+                   external_data: Optional[ExternalData] = None) -> io.IOBase:
+        """
+       Query method that returns the result as an io.IOBase iterator
+       :param query: Query statement/format string
+       :param parameters: Optional dictionary used to format the query
+       :param settings: Optional dictionary of ClickHouse settings (key/string values)
+       :param fmt: ClickHouse output format
+       :param use_database  Send the database parameter to ClickHouse so the command will be executed in the client
+        database context.
+       :param external_data  External data to send with the query
+       :return: io.IOBase stream/iterator for the result
+       """
+
     # pylint: disable=duplicate-code,too-many-arguments,unused-argument
     def query_np(self,
                  query: Optional[str] = None,
@@ -487,12 +505,11 @@ class Client(ABC):
         :return: Generator that yields a PyArrow.Table for per block representing the result set
         """
         settings = self._update_arrow_settings(settings, use_strings)
-        return to_arrow_batches(self.raw_query(query,
-                                               parameters,
-                                               settings,
-                                               fmt='ArrowStream',
-                                               external_data=external_data,
-                                               stream=True))
+        return to_arrow_batches(self.raw_stream(query,
+                                                parameters,
+                                                settings,
+                                                fmt='ArrowStream',
+                                                external_data=external_data))
 
     def _update_arrow_settings(self,
                                settings: Optional[Dict[str, Any]],
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py
index 1a35470b437..7202bc2ef51 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/httpclient.py
@@ -1,3 +1,4 @@
+import io
 import json
 import logging
 import re
@@ -436,27 +437,36 @@ class HttpClient(Client):
             else:
                 self._error_handler(response)
 
-    def ping(self):
-        """
-        See BaseClient doc_string for this method
-        """
-        try:
-            response = self.http.request('GET', f'{self.url}/ping', timeout=3)
-            return 200 <= response.status < 300
-        except HTTPError:
-            logger.debug('ping failed', exc_info=True)
-            return False
-
     def raw_query(self, query: str,
                   parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
                   settings: Optional[Dict[str, Any]] = None,
                   fmt: str = None,
                   use_database: bool = True,
-                  external_data: Optional[ExternalData] = None,
-                  stream: bool = False) -> Union[bytes, HTTPResponse]:
+                  external_data: Optional[ExternalData] = None) -> bytes:
         """
         See BaseClient doc_string for this method
         """
+        body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
+        return self._raw_request(body, params, fields=fields).data
+
+    def raw_stream(self, query: str,
+                   parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
+                   settings: Optional[Dict[str, Any]] = None,
+                   fmt: str = None,
+                   use_database: bool = True,
+                   external_data: Optional[ExternalData] = None) -> io.IOBase:
+        """
+        See BaseClient doc_string for this method
+        """
+        body, params, fields = self._prep_raw_query(query, parameters, settings, fmt, use_database, external_data)
+        return self._raw_request(body, params, fields=fields, stream=True)
+
+    def _prep_raw_query(self, query: str,
+                        parameters: Optional[Union[Sequence, Dict[str, Any]]],
+                        settings: Optional[Dict[str, Any]],
+                        fmt: str,
+                        use_database: bool,
+                        external_data: Optional[ExternalData]):
         final_query, bind_params = bind_query(query, parameters, self.server_tz)
         if fmt:
             final_query += f'\n FORMAT {fmt}'
@@ -472,8 +482,18 @@ class HttpClient(Client):
         else:
             body = final_query
             fields = None
-        response = self._raw_request(body, params, fields=fields, stream=stream)
-        return response if stream else response.data
+        return body, params, fields
+
+    def ping(self):
+        """
+        See BaseClient doc_string for this method
+        """
+        try:
+            response = self.http.request('GET', f'{self.url}/ping', timeout=3)
+            return 200 <= response.status < 300
+        except HTTPError:
+            logger.debug('ping failed', exc_info=True)
+            return False
 
     def close(self):
         if self._owns_pool_manager:
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py
index 549dfc37177..42957d8eb8d 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/driver/query.py
@@ -399,13 +399,13 @@ def format_query_value(value: Any, server_tz: tzinfo = pytz.UTC):
     if isinstance(value, date):
         return f"'{value.isoformat()}'"
     if isinstance(value, list):
-        return f"[{', '.join(format_query_value(x, server_tz) for x in value)}]"
+        return f"[{', '.join(str_query_value(x, server_tz) for x in value)}]"
     if isinstance(value, tuple):
-        return f"({', '.join(format_query_value(x, server_tz) for x in value)})"
+        return f"({', '.join(str_query_value(x, server_tz) for x in value)})"
     if isinstance(value, dict):
         if common.get_setting('dict_parameter_format') == 'json':
             return format_str(any_to_json(value).decode())
-        pairs = [format_query_value(k, server_tz) + ':' + format_query_value(v, server_tz)
+        pairs = [str_query_value(k, server_tz) + ':' + str_query_value(v, server_tz)
                  for k, v in value.items()]
         return f"{{{', '.join(pairs)}}}"
     if isinstance(value, Enum):
@@ -415,6 +415,10 @@ def format_query_value(value: Any, server_tz: tzinfo = pytz.UTC):
     return value
 
 
+def str_query_value(value: Any, server_tz: tzinfo = pytz.UTC):
+    return str(format_query_value(value, server_tz))
+
+
 # pylint: disable=too-many-branches
 def format_bind_value(value: Any, server_tz: tzinfo = pytz.UTC, top_level: bool = True):
     """
diff --git a/contrib/python/clickhouse-connect/clickhouse_connect/tools/testing.py b/contrib/python/clickhouse-connect/clickhouse_connect/tools/testing.py
index ef3f835e185..7084c71a406 100644
--- a/contrib/python/clickhouse-connect/clickhouse_connect/tools/testing.py
+++ b/contrib/python/clickhouse-connect/clickhouse_connect/tools/testing.py
@@ -1,7 +1,7 @@
 from typing import Sequence, Optional, Union, Dict, Any
 
 from clickhouse_connect.driver import Client
-from clickhouse_connect.driver.query import format_query_value, quote_identifier
+from clickhouse_connect.driver.query import quote_identifier, str_query_value
 
 
 class TableContext:
@@ -44,8 +44,7 @@ class TableContext:
         if self.settings:
             create_cmd += ' SETTINGS '
             for key, value in self.settings.items():
-
-                create_cmd += f'{key} = {format_query_value(value)}, '
+                create_cmd += f'{key} = {str_query_value(value)}, '
             if create_cmd.endswith(', '):
                 create_cmd = create_cmd[:-2]
         self.client.command(create_cmd)
diff --git a/contrib/python/clickhouse-connect/ya.make b/contrib/python/clickhouse-connect/ya.make
index 9c9ea8d7dfc..c48e09768d4 100644
--- a/contrib/python/clickhouse-connect/ya.make
+++ b/contrib/python/clickhouse-connect/ya.make
@@ -2,7 +2,7 @@
 
 PY3_LIBRARY()
 
-VERSION(0.7.6)
+VERSION(0.7.7)
 
 LICENSE(Apache-2.0)