diff --git a/CHANGELOG.md b/CHANGELOG.md index e8134ce98..6668aa8d2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/databricks/sqlalchemy/_parse.py b/src/databricks/sqlalchemy/_parse.py index 3d96e1603..42c4774d7 100644 --- a/src/databricks/sqlalchemy/_parse.py +++ b/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, } diff --git a/src/databricks/sqlalchemy/test_local/e2e/test_basic.py b/src/databricks/sqlalchemy/test_local/e2e/test_basic.py index 7d3cad511..3696356c9 100644 --- a/src/databricks/sqlalchemy/test_local/e2e/test_basic.py +++ b/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