From 185f196d269643c5a54faf1a38c86e32e144ae56 Mon Sep 17 00:00:00 2001
From: Konstantin Lebedev <kostyan.lebedev@gmail.com>
Date: Wed, 1 Feb 2023 22:05:14 +0300
Subject: [PATCH] Pass timezone to DateTime #230

---
 .../drivers/compilers/typecompiler.py         |  7 +++--
 clickhouse_sqlalchemy/types/common.py         |  9 ++++---
 tests/types/test_datetime.py                  | 26 ++++++++++++++-----
 3 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/clickhouse_sqlalchemy/drivers/compilers/typecompiler.py b/clickhouse_sqlalchemy/drivers/compilers/typecompiler.py
index 02abaecb..26647841 100644
--- 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
index df02186c..d5659f7a 100644
--- 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
index 8c6354c2..9a7862ec 100644
--- 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):