Deal with datetime objects (#327)

* Deal with datetime objects

* Allow changing milli/microseconds precision

* Reference self instead of reinstantiating ParamEscaper

Author: bjgbeelen

pyhive/common.py

@@ -14,6 +14,7 @@
 import abc
 import collections
 import time
+import datetime
 from future.utils import with_metaclass
 from itertools import islice
 
@@ -201,6 +202,10 @@ def __cmp__(self, other):
 
 
 class ParamEscaper(object):
+    _DATE_FORMAT = "%Y-%m-%d"
+    _TIME_FORMAT = "%H:%M:%S.%f"
+    _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT)
+
     def escape_args(self, parameters):
         if isinstance(parameters, dict):
             return {k: self.escape_item(v) for k, v in parameters.items()}
@@ -228,6 +233,11 @@ def escape_sequence(self, item):
         l = map(str, map(self.escape_item, item))
         return '(' + ','.join(l) + ')'
 
+    def escape_datetime(self, item, format, cutoff=0):
+        dt_str = item.strftime(format)
+        formatted = dt_str[:-cutoff] if cutoff and format.endswith(".%f") else dt_str
+        return "'{}'".format(formatted)
+
     def escape_item(self, item):
         if item is None:
             return 'NULL'
@@ -237,6 +247,10 @@ def escape_item(self, item):
             return self.escape_string(item)
         elif isinstance(item, collections.Iterable):
             return self.escape_sequence(item)
+        elif isinstance(item, datetime.datetime):
+            return self.escape_datetime(item, self._DATETIME_FORMAT)
+        elif isinstance(item, datetime.date):
+            return self.escape_datetime(item, self._DATE_FORMAT)
         else:
             raise exc.ProgrammingError("Unsupported object {}".format(item))
 

pyhive/presto.py

@@ -15,6 +15,7 @@
 from pyhive.exc import *  # noqa
 import base64
 import getpass
+import datetime
 import logging
 import requests
 from requests.auth import HTTPBasicAuth
@@ -32,7 +33,16 @@
 paramstyle = 'pyformat'  # Python extended format codes, e.g. ...WHERE name=%(name)s
 
 _logger = logging.getLogger(__name__)
-_escaper = common.ParamEscaper()
+
+
+class PrestoParamEscaper(common.ParamEscaper):
+    def escape_datetime(self, item, format):
+        _type = "timestamp" if isinstance(item, datetime.datetime) else "date"
+        formatted = super(PrestoParamEscaper, self).escape_datetime(item, format, 3)
+        return "{} {}".format(_type, formatted)
+
+
+_escaper = PrestoParamEscaper()
 
 
 def connect(*args, **kwargs):

pyhive/tests/test_common.py

@@ -2,7 +2,7 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
 from pyhive import common
-
+import datetime
 import unittest
 
 
@@ -34,3 +34,7 @@ def test_escape_args(self):
                          ("('a','b','c')",))
         self.assertEqual(escaper.escape_args((['你好', 'b', 'c'],)),
                          ("('你好','b','c')",))
+        self.assertEqual(escaper.escape_args((datetime.date(2020, 4, 17),)),
+                         ("'2020-04-17'",))
+        self.assertEqual(escaper.escape_args((datetime.datetime(2020, 4, 17, 12, 0, 0, 123456),)),
+                         ("'2020-04-17 12:00:00.123456'",))

pyhive/tests/test_presto.py

@@ -17,6 +17,7 @@
 from pyhive.tests.dbapi_test_case import with_cursor
 import mock
 import unittest
+import datetime
 
 _HOST = 'localhost'
 _PORT = '8080'
@@ -32,6 +33,14 @@ def test_bad_protocol(self):
         self.assertRaisesRegexp(ValueError, 'Protocol must be',
                                 lambda: presto.connect('localhost', protocol='nonsense').cursor())
 
+    def test_escape_args(self):
+        escaper = presto.PrestoParamEscaper()
+
+        self.assertEqual(escaper.escape_args((datetime.date(2020, 4, 17),)),
+                         ("date '2020-04-17'",))
+        self.assertEqual(escaper.escape_args((datetime.datetime(2020, 4, 17, 12, 0, 0, 123456),)),
+                         ("timestamp '2020-04-17 12:00:00.123'",))
+
     @with_cursor
     def test_description(self, cursor):
         cursor.execute('SELECT 1 AS foobar FROM one_row')