Fix DateTime and DateTime64 reflection #148

clickhouse_sqlalchemy/drivers/base.py

@@ -13,6 +13,7 @@
 from .compilers.sqlcompiler import ClickHouseSQLCompiler
 from .compilers.typecompiler import ClickHouseTypeCompiler
 from .reflection import ClickHouseInspector
+from .util import get_inner_spec
 from .. import types
 
 # Column specifications
@@ -263,9 +264,15 @@ def _get_column_type(self, name, spec):
             type_enum = enum.Enum('%s_enum' % name, options)
             return lambda: coltype(type_enum)
 
-        elif spec.lower().startswith('decimal'):
+        elif spec.startswith('Decimal'):
             coltype = self.ischema_names['Decimal']
             return coltype(*self._parse_decimal_params(spec))
+        elif spec.startswith('DateTime64'):
+            coltype = self.ischema_names['DateTime64']
+            return coltype(*self._parse_detetime64_params(spec))
+        elif spec.startswith('DateTime'):
+            coltype = self.ischema_names['DateTime']
+            return coltype(*self._parse_detetime_params(spec))
         else:
             try:
                 return self.ischema_names[spec]
@@ -276,10 +283,28 @@ def _get_column_type(self, name, spec):
 
     @staticmethod
     def _parse_decimal_params(spec):
-        ints = spec.split('(')[-1].split(')')[0]  # get all data in brackets
-        precision, scale = ints.split(',')
+        inner_spec = get_inner_spec(spec)
+        precision, scale = inner_spec.split(',')
         return int(precision.strip()), int(scale.strip())
 
+    @staticmethod
+    def _parse_detetime64_params(spec):
+        inner_spec = get_inner_spec(spec)
+        if not inner_spec:
+            return []
+        params = inner_spec.split(',', 1)
+        params[0] = int(params[0])
+        if len(params) > 1:
+            params[1] = params[1].strip()
+        return params
+
+    @staticmethod
+    def _parse_detetime_params(spec):
+        inner_spec = get_inner_spec(spec)
+        if not inner_spec:
+            return []
+        return [inner_spec]
+
     @staticmethod
     def _parse_options(option_string):
         options = dict()

clickhouse_sqlalchemy/drivers/util.py

@@ -0,0 +1,18 @@
+
+def get_inner_spec(spec):
+    brackets = 0
+    offset = spec.find('(')
+    if offset == -1:
+        return ''
+    i = offset
+    for i, ch in enumerate(spec[offset:], offset):
+        if ch == '(':
+            brackets += 1
+
+        elif ch == ')':
+            brackets -= 1
+
+        if brackets == 0:
+            break
+
+    return spec[offset + 1:i]

tests/test_reflection.py

@@ -133,3 +133,34 @@ def test_decimal(self):
         self.assertIsInstance(coltype, types.Decimal)
         self.assertEqual(coltype.precision, 38)
         self.assertEqual(coltype.scale, 38)
+    def test_datetime64(self):
+        coltype = self._type_round_trip(
+            types.DateTime64(2, 'Europe/Moscow')
+        )[0]['type']
+
+        self.assertIsInstance(coltype, types.DateTime64)
+        self.assertEqual(coltype.precision, 2)
+        self.assertEqual(coltype.timezone, "'Europe/Moscow'")
+
+        coltype = self._type_round_trip(
+            types.DateTime64()
+        )[0]['type']
+
+        self.assertIsInstance(coltype, types.DateTime64)
+        self.assertEqual(coltype.precision, 3)
+        self.assertIsNone(coltype.timezone)
+
+    def test_datetime(self):
+        coltype = self._type_round_trip(
+            types.DateTime('Europe/Moscow')
+        )[0]['type']
+
+        self.assertIsInstance(coltype, types.DateTime)
+        self.assertEqual(coltype.timezone, "'Europe/Moscow'")
+
+        coltype = self._type_round_trip(
+            types.DateTime()
+        )[0]['type']
+
+        self.assertIsInstance(coltype, types.DateTime)
+        self.assertIsNone(coltype.timezone)