Register date_trunc function for crate dialect with correct return type

Author: autophagy

CHANGES.txt

@@ -5,6 +5,9 @@ Changes for crate
 Unreleased
 ==========
 
+- Made it so that the SQLAlchemy dialect is now aware of the return type of the
+  ``date_trunc`` function.
+
 2019/09/19 0.23.2
 =================
 

src/crate/client/sqlalchemy/dialect.py

@@ -24,6 +24,7 @@
 
 from sqlalchemy import types as sqltypes
 from sqlalchemy.engine import default, reflection
+from sqlalchemy.sql import functions
 
 from .compiler import (
     CrateCompiler,
@@ -306,3 +307,8 @@ def _create_column_info(self, row):
 
     def _resolve_type(self, type_):
         return TYPES_MAP.get(type_, sqltypes.UserDefinedType)
+
+
+class DateTrunc(functions.GenericFunction):
+    name = "date_trunc"
+    type = sqltypes.TIMESTAMP

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

@@ -12,6 +12,7 @@
 from .create_table_test import CreateTableTest
 from .array_test import SqlAlchemyArrayTypeTest
 from .dialect_test import DialectTest
+from .function_test import FunctionTest
 from ..sa_version import SA_1_1, SA_VERSION
 
 
@@ -28,6 +29,7 @@ def test_suite():
     tests.addTest(makeSuite(SqlAlchemyInsertFromSelectTest))
     tests.addTest(makeSuite(SqlAlchemyInsertFromSelectTest))
     tests.addTest(makeSuite(DialectTest))
+    tests.addTest(makeSuite(FunctionTest))
     if SA_VERSION >= SA_1_1:
         tests.addTest(makeSuite(SqlAlchemyArrayTypeTest))
     return tests

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

@@ -0,0 +1,44 @@
+# -*- coding: utf-8; -*-
+#
+# Licensed to CRATE Technology GmbH ("Crate") under one or more contributor
+# license agreements.  See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.  Crate licenses
+# this file to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.  You may
+# obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+# License for the specific language governing permissions and limitations
+# under the License.
+#
+# However, if you have executed another commercial license agreement
+# with Crate these terms will supersede the license and you may use the
+# software solely pursuant to the terms of the relevant commercial agreement.
+
+from unittest import TestCase
+
+import sqlalchemy as sa
+from sqlalchemy.sql.sqltypes import TIMESTAMP
+from sqlalchemy.ext.declarative import declarative_base
+
+
+class FunctionTest(TestCase):
+    def setUp(self):
+        Base = declarative_base(bind=sa.create_engine("crate://"))
+
+        class Character(Base):
+            __tablename__ = "characters"
+            name = sa.Column(sa.String, primary_key=True)
+            timestamp = sa.Column(sa.DateTime)
+
+        self.Character = Character
+
+    def test_date_trunc_type_is_timestamp(self):
+        f = sa.func.date_trunc("minute", self.Character.timestamp)
+        self.assertEqual(len(f.base_columns), 1)
+        for col in f.base_columns:
+            self.assertIsInstance(col.type, TIMESTAMP)