Add proper timestamp timezone handling

tzonghao · tzonghao · commit cbde490024e4 · 2016-09-09T16:33:31.000-04:00
Fix bug influxdata#250
diff --git a/influxdb/line_protocol.py b/influxdb/line_protocol.py
@@ -5,25 +5,26 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-from calendar import timegm
 from copy import copy
 from datetime import datetime
 from numbers import Integral
 
+from pytz import UTC
 from dateutil.parser import parse
 from six import binary_type, text_type, integer_types, PY2
 
+EPOCH = UTC.localize(datetime.utcfromtimestamp(0))
+
 
 def _convert_timestamp(timestamp, precision=None):
     if isinstance(timestamp, Integral):
         return timestamp  # assume precision is correct if timestamp is int
     if isinstance(_get_unicode(timestamp), text_type):
         timestamp = parse(timestamp)
     if isinstance(timestamp, datetime):
-        ns = (
-            timegm(timestamp.utctimetuple()) * 1e9 +
-            timestamp.microsecond * 1e3
-        )
+        if not timestamp.tzinfo:
+            timestamp = UTC.localize(timestamp)
+        ns = (timestamp - EPOCH).total_seconds() * 1e9
         if precision is None or precision == 'n':
             return ns
         elif precision == 'u':
@@ -36,7 +37,6 @@ def _convert_timestamp(timestamp, precision=None):
             return ns / 1e9 / 60
         elif precision == 'h':
             return ns / 1e9 / 3600
-
     raise ValueError(timestamp)
 
 
diff --git a/influxdb/tests/test_line_protocol.py b/influxdb/tests/test_line_protocol.py
@@ -5,7 +5,9 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
+from datetime import datetime
 import unittest
+from pytz import UTC, timezone
 
 from influxdb import line_protocol
 
@@ -40,6 +42,35 @@ def test_make_lines(self):
             'bool_val=True,float_val=1.1,int_val=1i,string_val="hello!"\n'
         )
 
+    def test_timezone(self):
+        dt = datetime(2009, 11, 10, 23, 0, 0, 123456)
+        utc = UTC.localize(dt)
+        berlin = timezone('Europe/Berlin').localize(dt)
+        eastern = berlin.astimezone(timezone('US/Eastern'))
+        data = {
+            "points": [
+                {"measurement": "A", "fields": {"val": 1},
+                 "time": 0},
+                {"measurement": "A", "fields": {"val": 1},
+                 "time": "2009-11-10T23:00:00.123456Z"},
+                {"measurement": "A", "fields": {"val": 1}, "time": dt},
+                {"measurement": "A", "fields": {"val": 1}, "time": utc},
+                {"measurement": "A", "fields": {"val": 1}, "time": berlin},
+                {"measurement": "A", "fields": {"val": 1}, "time": eastern},
+            ]
+        }
+        self.assertEqual(
+            line_protocol.make_lines(data),
+            '\n'.join([
+                'A val=1i 0',
+                'A val=1i 1257894000123456000',
+                'A val=1i 1257894000123456000',
+                'A val=1i 1257894000123456000',
+                'A val=1i 1257890400123456000',
+                'A val=1i 1257890400123456000',
+            ]) + '\n'
+        )
+
     def test_string_val_newline(self):
         data = {
             "points": [
