TIMESTAMP_NTZ datatype used by DataFrame.to_sql but fails on read_sql_table

[cdcadman]

I'm using:
python 3.12
pandas==2.1.3
sqlalchemy==2.0.22
databricks-sql-connector==3.0.0
Windows

If I have a pandas DataFrame with a datetime64[ns] data type and use to_sql to write it into databricks, then read_sql_query to read the same table, it fails to handle the TIMESTAMP_NTZ data type.

Here is example code:

from os import environ

from pandas import DataFrame, Timestamp, read_sql_table
from sqlalchemy import create_engine
from sqlalchemy.event import listen

host = environ["DATABRICKS_HOST"]
http_path = environ["DATABRICKS_HTTP_PATH"] # SQL Warehouse
catalog = environ["DATABRICKS_CATALOG"]
schema = environ["DATABRICKS_SCHEMA"]


def do_connect(dialect, conn_rec, cargs, cparams):
 cparams["access_token"] = environ["DATABRICKS_ACCESS_TOKEN"]


engine = create_engine(
 f"databricks://{host}:443?http_path={http_path}&catalog={catalog}&schema={schema}"
)
listen(engine, "do_connect", do_connect)
with engine.connect() as conn:
 dtf = DataFrame(
 {
 "timestamp_field": [
 Timestamp("12/31/2023 12:00pm"),
 Timestamp("1/31/2023 2:00pm"),
 ]
 }
 )
 conn.exec_driver_sql("drop table if exists test_timestamp")
 dtf.to_sql("test_timestamp", conn, index=False)
 read_sql_table("test_timestamp", conn)

Here is my output:

C:\Temp>python databricks_timestamp_ntz.py
C:\Temp\databricks_timestamp_ntz.py:17: SADeprecationWarning: The dbapi() classmethod on dialect classes has been renamed to import_dbapi(). Implement an import_dbapi() classmethod directly on class <class 'databricks.sqlalchemy.base.DatabricksDialect'> to remove this warning; the old .dbapi() classmethod may be maintained for backwards compatibility.
 engine = create_engine(
Traceback (most recent call last):
 File "C:\Temp\databricks_timestamp_ntz.py", line 32, in <module>
 read_sql_table("test_timestamp", conn)
 File "C:\Program Files\Python312\Lib\site-packages\pandas\io\sql.py", line 348, in read_sql_table
 table = pandas_sql.read_table(
 ^^^^^^^^^^^^^^^^^^^^^^
 File "C:\Program Files\Python312\Lib\site-packages\pandas\io\sql.py", line 1665, in read_table
 self.meta.reflect(bind=self.con, only=[table_name], views=True)
 File "C:\Program Files\Python312\Lib\site-packages\sqlalchemy\sql\schema.py", line 5788, in reflect
 _reflect_info = insp._get_reflection_info(
 ^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "C:\Program Files\Python312\Lib\site-packages\sqlalchemy\engine\reflection.py", line 2006, in _get_reflection_info
 columns=run(
 ^^^^
 File "C:\Program Files\Python312\Lib\site-packages\sqlalchemy\engine\reflection.py", line 1992, in run
 res = meth(filter_names=_fn, **kw)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "C:\Program Files\Python312\Lib\site-packages\sqlalchemy\engine\reflection.py", line 921, in get_multi_columns
 table_col_defs = dict(
 ^^^^^
 File "C:\Program Files\Python312\Lib\site-packages\sqlalchemy\engine\default.py", line 1099, in _default_multi_reflect
 single_tbl_method(
 File "C:\Program Files\Python312\Lib\site-packages\databricks\sqlalchemy\base.py", line 150, in get_columns
 row_dict = parse_column_info_from_tgetcolumnsresponse(col)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 File "C:\Program Files\Python312\Lib\site-packages\databricks\sqlalchemy\_parse.py", line 330, in parse_column_info_from_tgetcolumnsresponse
 _col_type = GET_COLUMNS_TYPE_MAP[_raw_col_type]
 ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
KeyError: 'timestamp_ntz'

[susodapop]

Thanks for the report. This happens because TGetColumnsResponse from runtime returns an unexpected value. Fix incoming.

[susodapop]

PR with the fix is now open. Once it merges we'll cut a release that should get you back in business 👍

[susodapop]

You can now pip install databricks-sql-connector==3.0.1 which incorporates the fix.

[cdcadman]

@susodapop , thank you for quickly fixing this!