diff --git a/clickhouse_sqlalchemy/drivers/base.py b/clickhouse_sqlalchemy/drivers/base.py index b93c6aa0..b5d23fa2 100644 --- a/clickhouse_sqlalchemy/drivers/base.py +++ b/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() diff --git a/clickhouse_sqlalchemy/drivers/util.py b/clickhouse_sqlalchemy/drivers/util.py new file mode 100644 index 00000000..25783daf --- /dev/null +++ b/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] diff --git a/tests/test_reflection.py b/tests/test_reflection.py index 91c65013..e8d74f20 100644 --- a/tests/test_reflection.py +++ b/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)