Pass timezone to DateTime #230

xzkostyan · Feb 1, 2023 · 185f196 · 185f196
1 parent e2fe139
commit 185f196
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 12 deletions.
diff --git a/clickhouse_sqlalchemy/drivers/compilers/typecompiler.py b/clickhouse_sqlalchemy/drivers/compilers/typecompiler.py
@@ -61,11 +61,14 @@ def visit_date(self, type_, **kw):
         return 'Date'
 
     def visit_datetime(self, type_, **kw):
-        return 'DateTime'
+        if type_.timezone:
+            return "DateTime('%s')" % type_.timezone
+        else:
+            return 'DateTime'
 
     def visit_datetime64(self, type_, **kw):
         if type_.timezone:
-            return 'DateTime64(%s, \'%s\')' % (type_.precision, type_.timezone)
+            return "DateTime64(%s, '%s')" % (type_.precision, type_.timezone)
         else:
             return 'DateTime64(%s)' % type_.precision
 

diff --git a/clickhouse_sqlalchemy/types/common.py b/clickhouse_sqlalchemy/types/common.py
@@ -128,17 +128,20 @@ class Date(types.Date, ClickHouseTypeEngine):
     __visit_name__ = 'date'
 
 
-class DateTime(types.Date, ClickHouseTypeEngine):
+class DateTime(types.DateTime, ClickHouseTypeEngine):
     __visit_name__ = 'datetime'
 
+    def __init__(self, timezone=None):
+        super(DateTime, self).__init__()
+        self.timezone = timezone
+
 
 class DateTime64(DateTime, ClickHouseTypeEngine):
     __visit_name__ = 'datetime64'
 
     def __init__(self, precision=3, timezone=None):
         self.precision = precision
-        self.timezone = timezone
-        super(DateTime64, self).__init__()
+        super(DateTime64, self).__init__(timezone=timezone)
 
 
 class Enum(types.Enum, ClickHouseTypeEngine):

diff --git a/tests/types/test_datetime.py b/tests/types/test_datetime.py
@@ -9,18 +9,30 @@
 
 
 class DateTimeCompilationTestCase(CompilationTestCase):
-    table = Table(
-        'test', CompilationTestCase.metadata(),
-        Column('x', types.DateTime, primary_key=True),
-        engines.Memory()
-    )
-
     def test_create_table(self):
+        table = Table(
+            'test', CompilationTestCase.metadata(),
+            Column('x', types.DateTime, primary_key=True),
+            engines.Memory()
+        )
+
         self.assertEqual(
-            self.compile(CreateTable(self.table)),
+            self.compile(CreateTable(table)),
             'CREATE TABLE test (x DateTime) ENGINE = Memory'
         )
 
+    def test_create_table_with_timezone(self):
+        table = Table(
+            'test', CompilationTestCase.metadata(),
+            Column('x', types.DateTime('Europe/Moscow'), primary_key=True),
+            engines.Memory()
+        )
+
+        self.assertEqual(
+            self.compile(CreateTable(table)),
+            "CREATE TABLE test (x DateTime('Europe/Moscow')) ENGINE = Memory"
+        )
+
 
 @with_native_and_http_sessions
 class DateTimeTestCase(BaseTestCase):