cqlsh: Add local timezone support to cqlsh

patch by Stefan Podkowinski; reviewed by Paulo Motta for CASSANDRA-10397

CHANGES.txt

@@ -19,6 +19,7 @@ Merged from 2.1:
  * Avoid major compaction mixing repaired and unrepaired sstables in DTCS (CASSANDRA-11113)
  * Make it clear what DTCS timestamp_resolution is used for (CASSANDRA-11041)
  * test_bulk_round_trip_blogposts is failing occasionally (CASSANDRA-10938)
+ * (cqlsh) Support timezone conversion using pytz (CASSANDRA-10397)
 
 
 2.2.5

bin/cqlsh.py

@@ -666,6 +666,7 @@ class Shell(cmd.Cmd):
                  display_timestamp_format=DEFAULT_TIMESTAMP_FORMAT,
                  display_date_format=DEFAULT_DATE_FORMAT,
                  display_float_precision=DEFAULT_FLOAT_PRECISION,
+                 display_timezone=None,
                  max_trace_wait=DEFAULT_MAX_TRACE_WAIT,
                  ssl=False,
                  single_statement=None,
@@ -715,6 +716,8 @@ class Shell(cmd.Cmd):
 
         self.display_float_precision = display_float_precision
 
+        self.display_timezone = display_timezone
+
         # If there is no schema metadata present (due to a schema mismatch), force schema refresh
         if not self.conn.metadata.keyspaces:
             self.refresh_schema_metadata_best_effort()
@@ -801,7 +804,8 @@ class Shell(cmd.Cmd):
             self.decoding_errors.append(val)
         try:
             dtformats = DateTimeFormat(timestamp_format=self.display_timestamp_format,
-                                       date_format=self.display_date_format, nanotime_format=self.display_nanotime_format)
+                                       date_format=self.display_date_format, nanotime_format=self.display_nanotime_format,
+                                       timezone=self.display_timezone)
             return format_value(val, self.output_codec.name,
                                 addcolor=self.color, date_time_format=dtformats,
                                 float_precision=self.display_float_precision, **kwargs)
@@ -2349,6 +2353,7 @@ def read_options(cmdlineargs, environment):
     optvalues.field_size_limit = option_with_default(configs.getint, 'csv', 'field_size_limit', csv.field_size_limit())
     optvalues.max_trace_wait = option_with_default(configs.getfloat, 'tracing', 'max_trace_wait',
                                                    DEFAULT_MAX_TRACE_WAIT)
+    optvalues.timezone = option_with_default(configs.get, 'ui', 'timezone', None)
 
     optvalues.debug = False
     optvalues.file = None
@@ -2467,6 +2472,36 @@ def main(options, hostname, port):
         sys.stderr.write("Using CQL driver: %s\n" % (cassandra,))
         sys.stderr.write("Using connect timeout: %s seconds\n" % (options.connect_timeout,))
 
+    # create timezone based on settings, environment or auto-detection
+    timezone = None
+    if options.timezone or 'TZ' in os.environ:
+        try:
+            import pytz
+            if options.timezone:
+                try:
+                    timezone = pytz.timezone(options.timezone)
+                except:
+                    sys.stderr.write("Warning: could not recognize timezone '%s' specified in cqlshrc\n\n" % (options.timezone))
+            if 'TZ' in os.environ:
+                try:
+                    timezone = pytz.timezone(os.environ['TZ'])
+                except:
+                    sys.stderr.write("Warning: could not recognize timezone '%s' from environment value TZ\n\n" % (os.environ['TZ']))
+        except ImportError:
+            sys.stderr.write("Warning: Timezone defined and 'pytz' module for timezone conversion not installed. Timestamps will be displayed in UTC timezone.\n\n")
+
+    # try auto-detect timezone if tzlocal is installed
+    if not timezone:
+        try:
+            from tzlocal import get_localzone
+            timezone = get_localzone()
+        except ImportError:
+            # we silently ignore and fallback to UTC unless a custom timestamp format (which likely
+            # does contain a TZ part) was specified
+            if options.time_format != DEFAULT_TIMESTAMP_FORMAT:
+                sys.stderr.write("Warning: custom timestamp format specified in cqlshrc, but local timezone could not be detected.\n" +
+                                 "Either install Python 'tzlocal' module for auto-detection or specify client timezone in your cqlshrc.\n\n")
+
     try:
         shell = Shell(hostname,
                       port,
@@ -2483,6 +2518,7 @@ def main(options, hostname, port):
                       display_nanotime_format=options.nanotime_format,
                       display_date_format=options.date_format,
                       display_float_precision=options.float_precision,
+                      display_timezone=timezone,
                       max_trace_wait=options.max_trace_wait,
                       ssl=options.ssl,
                       single_statement=options.execute,

conf/cqlshrc.sample

@@ -32,6 +32,9 @@
 ;; Used for displaying timestamps (and reading them with COPY)
 ; datetimeformat = %Y-%m-%d %H:%M:%S%z
 
+;; Display timezone
+;timezone = Etc/UTC
+
 ;; The number of digits displayed after the decimal point
 ;; (note that increasing this to large numbers can result in unusual values)
 ; float_precision = 5

pylib/cqlshlib/formatting.py

@@ -107,10 +107,12 @@ if platform.system() == 'Windows':
 
 class DateTimeFormat():
 
-    def __init__(self, timestamp_format=DEFAULT_TIMESTAMP_FORMAT, date_format=DEFAULT_DATE_FORMAT, nanotime_format=DEFAULT_NANOTIME_FORMAT):
+    def __init__(self, timestamp_format=DEFAULT_TIMESTAMP_FORMAT, date_format=DEFAULT_DATE_FORMAT,
+                 nanotime_format=DEFAULT_NANOTIME_FORMAT, timezone=None):
         self.timestamp_format = timestamp_format
         self.date_format = date_format
         self.nanotime_format = nanotime_format
+        self.timezone = timezone
 
 
 def format_value_default(val, colormap, **_):
@@ -234,15 +236,17 @@ formatter_for('int')(format_integer_type)
 
 @formatter_for('datetime')
 def format_value_timestamp(val, colormap, date_time_format, quote=False, **_):
-    bval = strftime(date_time_format.timestamp_format, calendar.timegm(val.utctimetuple()))
+    bval = strftime(date_time_format.timestamp_format, calendar.timegm(val.utctimetuple()), timezone=date_time_format.timezone)
     if quote:
         bval = "'%s'" % bval
     return colorme(bval, colormap, 'timestamp')
 
 
-def strftime(time_format, seconds):
-    tzless_dt = datetime_from_timestamp(seconds)
-    return tzless_dt.replace(tzinfo=UTC()).strftime(time_format)
+def strftime(time_format, seconds, timezone=None):
+    ret_dt = datetime_from_timestamp(seconds).replace(tzinfo=UTC())
+    if timezone:
+        ret_dt = ret_dt.astimezone(timezone)
+    return ret_dt.strftime(time_format)
 
 
 @formatter_for('Date')

pylib/cqlshlib/test/test_cqlsh_output.py

@@ -370,6 +370,24 @@ class TestCqlshOutput(BaseTestCase):
             nnnnnnnn
             """),
         ), env={'TZ': 'Etc/UTC'})
+        try:
+            import pytz  # test only if pytz is available on PYTHONPATH
+            self.assertQueriesGiveColoredOutput((
+                ('''select timestampcol from has_all_types where num = 0;''', """
+                 timestampcol
+                 MMMMMMMMMMMM
+                --------------------------
+
+                 2012-05-14 09:53:20-0300
+                 GGGGGGGGGGGGGGGGGGGGGGGG
+
+
+                (1 rows)
+                nnnnnnnn
+                """),
+            ), env={'TZ': 'America/Sao_Paulo'})
+        except ImportError:
+            pass
 
     def test_boolean_output(self):
         self.assertCqlverQueriesGiveColoredOutput((