[PECO-1297] sqlalchemy: fix: can't read columns for tables containing a TIMESTAMP_NTZ column (databricks#296)

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+## 3.0.1 (unreleased)
+
+- Fix: SQLAlchemy dialect could not reflect TIMESTAMP_NTZ columns (#296)
+
 ## 3.0.0 (2023-11-17)
 
 - Remove support for Python 3.7

src/databricks/sqlalchemy/_parse.py

@@ -5,6 +5,8 @@
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.engine.interfaces import ReflectedColumn
 
+from databricks.sqlalchemy import _types as type_overrides
+
 """
 This module contains helper functions that can parse the contents
 of metadata and exceptions received from DBR. These are mostly just
@@ -293,7 +295,8 @@ def get_pk_strings_from_dte_output(
     "struct": sqlalchemy.types.String,
     "uniontype": sqlalchemy.types.String,
     "decimal": sqlalchemy.types.Numeric,
-    "timestamp": sqlalchemy.types.DateTime,
+    "timestamp": type_overrides.TIMESTAMP,
+    "timestamp_ntz": type_overrides.TIMESTAMP_NTZ,
     "date": sqlalchemy.types.Date,
 }
 

src/databricks/sqlalchemy/test_local/e2e/test_basic.py

@@ -1,21 +1,24 @@
-import os, datetime, decimal
-import pytest
+import datetime
+import decimal
+import os
+from typing import Tuple, Union
 from unittest import skipIf
+
+import pytest
 from sqlalchemy import (
-    create_engine,
-    select,
-    insert,
     Column,
     MetaData,
     Table,
     Text,
+    create_engine,
+    insert,
+    select,
     text,
 )
-from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column
-from sqlalchemy.types import SMALLINT, Integer, BOOLEAN, String, DECIMAL, Date
 from sqlalchemy.engine import Engine
-
-from typing import Tuple, Union
+from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
+from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
 
 try:
     from sqlalchemy.orm import declarative_base
@@ -344,8 +347,6 @@ class Base(DeclarativeBase):
 def test_inspector_smoke_test(samples_engine: Engine):
     """It does not appear that 3L namespace is supported here"""
 
-    from sqlalchemy.engine.reflection import Inspector
-
     schema, table = "nyctaxi", "trips"
 
     try:
@@ -431,3 +432,34 @@ def get_conn_user_agent(conn):
     c2.close()
 
     assert same_ua, f"User agents didn't match \n {ua1} \n {ua2}"
+
+
+@pytest.fixture
+def sample_table(metadata_obj: MetaData, db_engine: Engine):
+    """This fixture creates a sample table and cleans it up after the test is complete."""
+    from databricks.sqlalchemy._parse import GET_COLUMNS_TYPE_MAP
+
+    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
+
+    args = [
+        Column(colname, coltype) for colname, coltype in GET_COLUMNS_TYPE_MAP.items()
+    ]
+
+    SampleTable = Table(table_name, metadata_obj, *args)
+
+    metadata_obj.create_all(db_engine)
+
+    yield table_name
+
+    metadata_obj.drop_all(db_engine)
+
+
+def test_get_columns(db_engine, sample_table: str):
+    """Created after PECO-1297 and Github Issue #295 to verify that get_columsn behaves like it should for all known SQLAlchemy types"""
+
+    inspector = Inspector.from_engine(db_engine)
+
+    # this raises an exception if `parse_column_info_from_tgetcolumnsresponse` fails a lookup
+    columns = inspector.get_columns(sample_table)
+
+    assert True