SNOW-227101: fix to datetime columns being created without timezone (…

…TIMESTAMP_NTZ).

Commit 28e8031 introduced a bug where all sqlalchemy.DateTime(timezone=True) columns were created as DATETIME in Snowflake which is an alias to TIMESTAMP_NTZ.
This commit reverts datetime logic back so if the timezone is defined in datetime column it will create TIMESTAMP_TZ and TIMESTAMP_NTZ otherwise.

(cherry picked from commit e180b23)

Author: 0lmer

Diff for: src/snowflake/sqlalchemy/base.py

@@ -634,10 +634,10 @@ def visit_BLOB(self, type_, **kw):
         return "BINARY"
 
     def visit_datetime(self, type_, **kw):
-        return "datetime"
+        return self.visit_TIMESTAMP(type_, **kw)
 
     def visit_DATETIME(self, type_, **kw):
-        return "DATETIME"
+        return self.visit_TIMESTAMP(type_, **kw)
 
     def visit_TIMESTAMP_NTZ(self, type_, **kw):
         return "TIMESTAMP_NTZ"
@@ -649,7 +649,14 @@ def visit_TIMESTAMP_LTZ(self, type_, **kw):
         return "TIMESTAMP_LTZ"
 
     def visit_TIMESTAMP(self, type_, **kw):
-        return "TIMESTAMP"
+        is_local = kw.get('is_local', False)
+        timezone = kw.get('timezone', type_.timezone)
+        return "TIMESTAMP %s%s" % (
+            (timezone and "WITH" or "WITHOUT") + (
+                    is_local and " LOCAL" or "") + " TIME ZONE",
+            "(%d)" % type_.precision if getattr(type_, 'precision',
+                                               None) is not None else ""
+        )
 
     def visit_GEOGRAPHY(self, type_, **kw):
         return "GEOGRAPHY"

Diff for: tests/test_timestamp.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 
 import pytz
-from sqlalchemy import Column, Integer, MetaData, Table
+from sqlalchemy import Column, Integer, MetaData, Table, DateTime
 from sqlalchemy.sql import select
 
 from snowflake.sqlalchemy import TIMESTAMP_LTZ, TIMESTAMP_NTZ, TIMESTAMP_TZ
@@ -27,6 +27,8 @@ def test_create_table_timestamp_datatypes(engine_testaccount):
         Column("tsntz", TIMESTAMP_NTZ),
         Column("tsltz", TIMESTAMP_LTZ),
         Column("tstz", TIMESTAMP_TZ),
+        Column("dttz", DateTime(timezone=True)),
+        Column("dtntz", DateTime(timezone=False)),
     )
     metadata.create_all(engine_testaccount)
     try:
@@ -48,6 +50,8 @@ def test_inspect_timestamp_datatypes(engine_testaccount):
         Column("tsntz", TIMESTAMP_NTZ),
         Column("tsltz", TIMESTAMP_LTZ),
         Column("tstz", TIMESTAMP_TZ),
+        Column("dttz", DateTime(timezone=True)),
+        Column("dtntz", DateTime(timezone=False)),
     )
     metadata.create_all(engine_testaccount)
     try:
@@ -79,5 +83,7 @@ def test_inspect_timestamp_datatypes(engine_testaccount):
                 assert rows[1] == current_utctime
                 assert rows[2] == current_localtime
                 assert rows[3] == current_localtime_with_other_tz
+                assert rows[4] == current_localtime_with_other_tz
+                assert rows[5] == current_utctime
     finally:
         test_timestamp.drop(engine_testaccount)