From b1c78f8fbe03d5576949f5ee5c8ab43bd72ac2a5 Mon Sep 17 00:00:00 2001
From: Geoff Genz <geoff@clickhouse.com>
Date: Thu, 30 Mar 2023 04:27:15 -0600
Subject: [PATCH] Same timezone fix (#158)

* Don't apply timezone if server/client match

* Tiny cleanup
---
 CHANGELOG.md                            |  7 +++++++
 clickhouse_connect/VERSION              |  2 +-
 clickhouse_connect/driver/client.py     | 10 ++++++----
 clickhouse_connect/driver/httpclient.py |  2 +-
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 125794f8..8630aa6d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,13 @@ in a future release.  Starting with 0.5.9 the driver now requests ClickHouse pro
 The secondary effect of the `send_progress` argument -- to set `wait_end_of_query=1` -- is now handled automatically based
 on whether the query is streaming or not.
 
+## 0.5.18, 2023-03-30
+### Performance Improvement
+- The server timezone will not be applied (and Python datetime types will be timezone naive) if the client and server timezones match
+and the `get_client` apply_server_timezone parameter is True (the default).  This improves performance where client and server
+have the same (non-UTC) timezone.  To override this behavior and always apply a server timezone to the result, use `apply_server_timezone='always'`.
+This should fix https://github.com/ClickHouse/clickhouse-connect/issues/157
+
 
 ## 0.5.17, 2023-03-26
 ### Timezone Improvements
diff --git a/clickhouse_connect/VERSION b/clickhouse_connect/VERSION
index 0d240c6b..d2d81b78 100644
--- a/clickhouse_connect/VERSION
+++ b/clickhouse_connect/VERSION
@@ -1 +1 @@
-0.5.17
\ No newline at end of file
+0.5.18
\ No newline at end of file
diff --git a/clickhouse_connect/driver/client.py b/clickhouse_connect/driver/client.py
index 2abd367a..75ca56cb 100644
--- a/clickhouse_connect/driver/client.py
+++ b/clickhouse_connect/driver/client.py
@@ -1,6 +1,6 @@
 import io
 import logging
-from datetime import tzinfo
+from datetime import tzinfo, datetime
 
 import pytz
 
@@ -12,7 +12,7 @@
 from clickhouse_connect.common import version
 from clickhouse_connect.datatypes.registry import get_from_name
 from clickhouse_connect.datatypes.base import ClickHouseType
-from clickhouse_connect.driver.common import dict_copy, StreamContext, coerce_int
+from clickhouse_connect.driver.common import dict_copy, StreamContext, coerce_int, coerce_bool
 from clickhouse_connect.driver.constants import CH_VERSION_WITH_PROTOCOL, PROTOCOL_VERSION_WITH_LOW_CARD
 from clickhouse_connect.driver.exceptions import ProgrammingError
 from clickhouse_connect.driver.external import ExternalData
@@ -42,7 +42,7 @@ def __init__(self,
                  uri: str,
                  query_retries: int,
                  server_host_name: Optional[str],
-                 apply_server_timezone: Optional[bool]):
+                 apply_server_timezone: Optional[Union[str, bool]]):
         """
         Shared initialization of ClickHouse Connect client
         :param database: database name
@@ -52,7 +52,6 @@ def __init__(self,
         self.query_limit = coerce_int(query_limit)
         self.query_retries = coerce_int(query_retries)
         self.server_host_name = server_host_name
-        self.apply_server_timezone = apply_server_timezone is True
         self.server_tz = pytz.UTC
         self.server_version, server_tz, self.database = \
             tuple(self.command('SELECT version(), timezone(), currentDatabase()', use_database=False))
@@ -60,6 +59,9 @@ def __init__(self,
             self.server_tz = pytz.timezone(server_tz)
         except UnknownTimeZoneError:
             logger.warning('Warning, server is using an unrecognized timezone %s, will use UTC default', server_tz)
+        offsets_differ = datetime.now().astimezone().utcoffset() != datetime.now(tz=self.server_tz).utcoffset()
+        self.apply_server_timezone = apply_server_timezone == 'always' or (
+                    coerce_bool(apply_server_timezone) and offsets_differ)
         readonly = 'readonly'
         if not self.min_version('19.17'):
             readonly = common.get_setting('readonly')
diff --git a/clickhouse_connect/driver/httpclient.py b/clickhouse_connect/driver/httpclient.py
index 5591455b..4ed87d0c 100644
--- a/clickhouse_connect/driver/httpclient.py
+++ b/clickhouse_connect/driver/httpclient.py
@@ -68,7 +68,7 @@ def __init__(self,
                  http_proxy: Optional[str] = None,
                  https_proxy: Optional[str] = None,
                  server_host_name: Optional[str] = None,
-                 apply_server_timezone: Optional[bool] = True):
+                 apply_server_timezone: Optional[Union[str, bool]] = True):
         """
         Create an HTTP ClickHouse Connect client
         See clickhouse_connect.get_client for parameters