Followup: Return TIMESTAMP columns as native Python datetime objects

src/crate/client/cursor.py

@@ -21,6 +21,7 @@
 
 from .exceptions import ProgrammingError
 import warnings
+from datetime import datetime
 
 
 class Cursor(object):
@@ -49,9 +50,47 @@ def execute(self, sql, parameters=None, bulk_parameters=None):
 
         self._result = self.connection.client.sql(sql, parameters,
                                                   bulk_parameters)
-        if "rows" in self._result:
+
+        if "rows" not in self._result:
+            return
+
+        if "col_types" in self._result:
+            self.rows = iter(self._transform_result_types())
+
+        else:
             self.rows = iter(self._result["rows"])
 
+    def _transform_result_types(self):
+        """
+        Generate row items with column values converted to their corresponding
+        native Python types, based on information from `col_types`.
+
+        Currently, only converting to native `datetime` objects is implemented.
+        """
+        datetime_column_types = [11, 15]
+        datetime_columns_mask = [
+            True if col_type in datetime_column_types else False
+            for col_type in self._result["col_types"]
+        ]
+        for row in self._result["rows"]:
+            yield list(self._transform_datetime_columns(row, iter(datetime_columns_mask)))
+
+    @staticmethod
+    def _transform_datetime_columns(row, column_flags):
+        """
+        Convert all designated columns to native Python `datetime` objects.
+        """
+        for value in row:
+            try:
+                flag = next(column_flags)
+            except StopIteration:
+                break
+
+            if flag and value is not None:
+                value = datetime.fromtimestamp(value / 1e3)
+
+            yield value
+
     def executemany(self, sql, seq_of_parameters):
         """
         Prepare a database operation (query or command) and then execute it

src/crate/client/doctests/client.txt

@@ -204,6 +204,11 @@ fetchall()
 Cursor Description
 ==================
 
+.. hidden: Set time zone to UTC to make naive datetime objects deterministic.
+
+    >>> previous_timezone = os.environ.get("TZ")
+    >>> os.environ["TZ"] = "UTC"
+
 The ``description`` property of the cursor returns a sequence of 7-item
 sequences containing the column name as first parameter. Just the name field is
 supported, all other fields are 'None'::
@@ -212,9 +217,9 @@ supported, all other fields are 'None'::
     >>> result = cursor.fetchone()
     >>> pprint(result)
     ['Aldebaran',
-     1658167836758,
-     1658167836758,
-     1658167836758,
+     datetime.datetime(2022, 7, 18, 18, 10, 36, 758000),
+     datetime.datetime(2022, 7, 18, 18, 10, 36, 758000),
+     datetime.datetime(2022, 7, 18, 18, 10, 36, 758000),
      None,
      None,
      'Star System',
@@ -239,6 +244,11 @@ supported, all other fields are 'None'::
      ('description', None, None, None, None, None, None),
      ('details', None, None, None, None, None, None))
 
+.. hidden: Restore time zone
+
+    >>> if previous_timezone: os.environ["TZ"] = previous_timezone
+
+
 Closing the Cursor
 ==================
 

src/crate/client/doctests/http.txt

@@ -69,7 +69,8 @@ Issue a select statement against our with test data pre-filled crate instance::
     >>> http_client = HttpClient(crate_host)
     >>> result = http_client.sql('select name from locations order by name')
     >>> pprint(result)
-    {'cols': ['name'],
+    {'col_types': [4],
+     'cols': ['name'],
      'duration': ...,
      'rowcount': 13,
      'rows': [['Aldebaran'],

src/crate/client/http.py

@@ -315,7 +315,7 @@ class Client(object):
     Crate connection client using CrateDB's HTTP API.
     """
 
-    SQL_PATH = '/_sql'
+    SQL_PATH = '/_sql?types=true'
     """Crate URI path for issuing SQL statements."""
 
     retry_interval = 30
@@ -385,7 +385,7 @@ def __init__(self,
 
         self.path = self.SQL_PATH
         if error_trace:
-            self.path += '?error_trace=true'
+            self.path += '&error_trace=true'
 
     def close(self):
         for server in self.server_pool.values():

src/crate/client/sqlalchemy/dialect.py

@@ -89,6 +89,8 @@ def result_processor(self, dialect, coltype):
         def process(value):
             if not value:
                 return
+            if isinstance(value, datetime):
+                return value.date()
             try:
                 return datetime.utcfromtimestamp(value / 1e3).date()
             except TypeError:
@@ -128,6 +130,8 @@ def result_processor(self, dialect, coltype):
         def process(value):
             if not value:
                 return
+            if isinstance(value, datetime):
+                return value
             try:
                 return datetime.utcfromtimestamp(value / 1e3)
             except TypeError:

src/crate/client/sqlalchemy/doctests/itests.txt

@@ -80,9 +80,6 @@ Date should have been set at the insert due to default value via python method::
     >>> dt.day == now.day
     True
 
-    >>> (now - location.datetime_tz).seconds < 4
-    True
-
 Verify the return type of date and datetime::
 
     >>> type(location.date)

src/crate/client/test_http.py

@@ -431,12 +431,13 @@ def test_params(self):
         client = Client(['127.0.0.1:4200'], error_trace=True)
         parsed = urlparse(client.path)
         params = parse_qs(parsed.query)
+        self.assertEqual(params["types"], ["true"])
         self.assertEqual(params["error_trace"], ["true"])
         client.close()
 
     def test_no_params(self):
         client = Client()
-        self.assertEqual(client.path, "/_sql")
+        self.assertEqual(client.path, "/_sql?types=true")
         client.close()
 
 

src/crate/client/tests.py

@@ -149,6 +149,7 @@ def refresh(table):
 
 
 def setUpWithCrateLayer(test):
+    test.globs['os'] = os
     test.globs['HttpClient'] = http.Client
     test.globs['crate_host'] = crate_host
     test.globs['pprint'] = pprint