Return TIMESTAMP columns as native Python datetime objects to improve support for Apache Superset

.github/workflows/main.yml

@@ -0,0 +1,56 @@
+name: Tests
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+  workflow_dispatch:
+  schedule:
+    - cron: '0 2 * * *'
+
+
+jobs:
+  test:
+    name: "Python: ${{ matrix.python-version }}
+     SQLA: ${{ matrix.sqla-version }}
+     CrateDB: ${{ matrix.crate-version }}
+     on ${{ matrix.os }}"
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        crate-version: [nightly]
+        os: [ubuntu-latest]
+        sqla-version: ['1.1.18', '1.2.19', '1.3.23']
+        python-version: [3.5, 3.6, 3.7, 3.8, 3.9]
+
+    steps:
+      - uses: actions/checkout@master
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python bootstrap.py
+
+          # replace SQLAlchemy version
+          sed -ir 's/SQLAlchemy.*/SQLAlchemy = ${{ matrix.sqla-version }}/g' versions.cfg
+
+          # replace CrateDB version
+          if [ ${{ matrix.crate-version }} = "nightly" ]; then
+            sed -ir 's/releases/releases\/nightly/g' base.cfg
+            sed -ir 's/crate_server.*/crate_server = latest/g' versions.cfg
+          else
+            sed -ir 's/crate-/crate_/g' base.cfg
+            sed -ir 's/crate_server.*/crate_server = ${{ matrix.crate-version }}/g' versions.cfg
+          fi
+
+          bin/buildout -n -c base.cfg
+
+      - name: Test
+        run: |
+          bin/flake8
+          bin/coverage run bin/test -vv1

src/crate/client/cursor.py

@@ -22,6 +22,7 @@
 from .exceptions import ProgrammingError
 from distutils.version import StrictVersion
 import warnings
+from datetime import datetime
 
 BULK_INSERT_MIN_VERSION = StrictVersion("0.42.0")
 
@@ -52,8 +53,45 @@ def execute(self, sql, parameters=None, bulk_parameters=None):
 
         self._result = self.connection.client.sql(sql, parameters,
                                                   bulk_parameters)
+
         if "rows" in self._result:
-            self.rows = iter(self._result["rows"])
+            transformed_result = False
+            if "col_types" in self._result:
+                transformed_result = True
+                self.rows = self.result_set_transformed()
+
+            if not transformed_result:
+                self.rows = iter(self._result["rows"])
+
+    def result_set_transformed(self):
+        """
+        Generator that iterates over each row from the result set
+        """
+        rows_to_convert = [True if col_type == 11 or col_type == 15 else False for col_type in
+                           self._result["col_types"]]
+        for row in self._result["rows"]:
+            gen_flags = (flag for flag in rows_to_convert)
+            yield [t_row for t_row in self._transform_date_columns(row, gen_flags)]
+
+    @staticmethod
+    def _transform_date_columns(row, gen_flags):
+        """
+        Generates iterates over each value from a row and converts timestamps to pandas TIMESTAMP
+        """
+        for value in row:
+            try:
+                flag = next(gen_flags)
+            except StopIteration:
+                break
+
+            if not flag or value is None:
+                yield value
+            else:
+                if value < 0:
+                    yield None
+                else:
+                    value = datetime.fromtimestamp(value / 1000)
+                    yield value
 
     def executemany(self, sql, seq_of_parameters):
         """

src/crate/client/doctests/client.txt

@@ -179,7 +179,7 @@ supported, all other fields are 'None'::
     >>> result = cursor.fetchone()
     >>> pprint(result)
     ['Aldebaran',
-     1373932800000,
+     datetime.datetime(2013, 7, 15, 18, 0),
      None,
      None,
      None,

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'
     """Crate URI path for issuing SQL statements."""
 
     retry_interval = 30

src/crate/client/sqlalchemy/dialect.py

@@ -43,6 +43,7 @@
     "smallint": sqltypes.SmallInteger,
     "timestamp": sqltypes.TIMESTAMP,
     "timestamp with time zone": sqltypes.TIMESTAMP,
+    "timestamp without time zone": sqltypes.TIMESTAMP,
     "object": Object,
     "integer": sqltypes.Integer,
     "long": sqltypes.NUMERIC,
@@ -64,6 +65,7 @@
     TYPES_MAP["smallint_array"] = ARRAY(sqltypes.SmallInteger)
     TYPES_MAP["timestamp_array"] = ARRAY(sqltypes.TIMESTAMP)
     TYPES_MAP["timestamp with time zone_array"] = ARRAY(sqltypes.TIMESTAMP)
+    TYPES_MAP["timestamp without time zone_array"] = ARRAY(sqltypes.TIMESTAMP)
     TYPES_MAP["long_array"] = ARRAY(sqltypes.NUMERIC)
     TYPES_MAP["bigint_array"] = ARRAY(sqltypes.NUMERIC)
     TYPES_MAP["double_array"] = ARRAY(sqltypes.DECIMAL)
@@ -75,7 +77,6 @@
 except Exception:
     pass
 
-
 log = logging.getLogger(__name__)
 
 
@@ -91,6 +92,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:
@@ -130,6 +133,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:
@@ -261,7 +266,7 @@ def get_pk_constraint(self, engine, table_name, schema=None, **kw):
 
             def result_fun(result):
                 rows = result.fetchall()
-                return set(map(lambda el: el[0], rows))
+                return list(set(map(lambda el: el[0], rows)))
         else:
             query = """SELECT constraint_name
                    FROM information_schema.table_constraints

src/crate/client/sqlalchemy/tests/dialect_test.py

@@ -67,7 +67,7 @@ def test_pks_are_retrieved_depending_on_version_set(self):
         self.engine.dialect.server_version_info = (0, 54, 0)
         fake_cursor.rowcount = 1
         fake_cursor.fetchone = MagicMock(return_value=[["id", "id2", "id3"]])
-        eq_(insp.get_pk_constraint("characters")['constrained_columns'], {"id", "id2", "id3"})
+        eq_(insp.get_pk_constraint("characters")['constrained_columns'], ["id", "id2", "id3"])
         fake_cursor.fetchone.assert_called_once_with()
         in_("information_schema.table_constraints", self.executed_statement)
 
@@ -76,7 +76,7 @@ def test_pks_are_retrieved_depending_on_version_set(self):
         self.engine.dialect.server_version_info = (2, 3, 0)
         fake_cursor.rowcount = 3
         fake_cursor.fetchall = MagicMock(return_value=[["id"], ["id2"], ["id3"]])
-        eq_(insp.get_pk_constraint("characters")['constrained_columns'], {"id", "id2", "id3"})
+        eq_(insp.get_pk_constraint("characters")['constrained_columns'], ["id", "id2", "id3"])
         fake_cursor.fetchall.assert_called_once_with()
         in_("information_schema.key_column_usage", self.executed_statement)
 

src/crate/client/test_http.py

@@ -428,12 +428,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["error_trace"], ["true"])
+        print(params)
+        self.assertEqual(params["types?error_trace"], ["true"])
         client.close()
 
     def test_no_params(self):
         client = Client()
-        self.assertEqual(client.path, "/_sql")
+        self.assertEqual(client.path, "/_sql?types")
         client.close()