Skip to content

Commit 2a68448

Browse files
author
Jesse
authored
[PECO-1297] sqlalchemy: fix: can't read columns for tables containing a TIMESTAMP_NTZ column (databricks#296)
Signed-off-by: Jesse Whitehouse <[email protected]>
1 parent c730648 commit 2a68448

File tree

3 files changed

+51
-12
lines changed

3 files changed

+51
-12
lines changed

CHANGELOG.md

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
# Release History
22

3+
## 3.0.1 (unreleased)
4+
5+
- Fix: SQLAlchemy dialect could not reflect TIMESTAMP_NTZ columns (#296)
6+
37
## 3.0.0 (2023-11-17)
48

59
- Remove support for Python 3.7

src/databricks/sqlalchemy/_parse.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
from sqlalchemy.engine import CursorResult
66
from sqlalchemy.engine.interfaces import ReflectedColumn
77

8+
from databricks.sqlalchemy import _types as type_overrides
9+
810
"""
911
This module contains helper functions that can parse the contents
1012
of metadata and exceptions received from DBR. These are mostly just
@@ -293,7 +295,8 @@ def get_pk_strings_from_dte_output(
293295
"struct": sqlalchemy.types.String,
294296
"uniontype": sqlalchemy.types.String,
295297
"decimal": sqlalchemy.types.Numeric,
296-
"timestamp": sqlalchemy.types.DateTime,
298+
"timestamp": type_overrides.TIMESTAMP,
299+
"timestamp_ntz": type_overrides.TIMESTAMP_NTZ,
297300
"date": sqlalchemy.types.Date,
298301
}
299302

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

Lines changed: 43 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,24 @@
1-
import os, datetime, decimal
2-
import pytest
1+
import datetime
2+
import decimal
3+
import os
4+
from typing import Tuple, Union
35
from unittest import skipIf
6+
7+
import pytest
48
from sqlalchemy import (
5-
create_engine,
6-
select,
7-
insert,
89
Column,
910
MetaData,
1011
Table,
1112
Text,
13+
create_engine,
14+
insert,
15+
select,
1216
text,
1317
)
14-
from sqlalchemy.orm import Session, DeclarativeBase, Mapped, mapped_column
15-
from sqlalchemy.types import SMALLINT, Integer, BOOLEAN, String, DECIMAL, Date
1618
from sqlalchemy.engine import Engine
17-
18-
from typing import Tuple, Union
19+
from sqlalchemy.engine.reflection import Inspector
20+
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
21+
from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
1922

2023
try:
2124
from sqlalchemy.orm import declarative_base
@@ -344,8 +347,6 @@ class Base(DeclarativeBase):
344347
def test_inspector_smoke_test(samples_engine: Engine):
345348
"""It does not appear that 3L namespace is supported here"""
346349

347-
from sqlalchemy.engine.reflection import Inspector
348-
349350
schema, table = "nyctaxi", "trips"
350351

351352
try:
@@ -431,3 +432,34 @@ def get_conn_user_agent(conn):
431432
c2.close()
432433

433434
assert same_ua, f"User agents didn't match \n {ua1} \n {ua2}"
435+
436+
437+
@pytest.fixture
438+
def sample_table(metadata_obj: MetaData, db_engine: Engine):
439+
"""This fixture creates a sample table and cleans it up after the test is complete."""
440+
from databricks.sqlalchemy._parse import GET_COLUMNS_TYPE_MAP
441+
442+
table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
443+
444+
args = [
445+
Column(colname, coltype) for colname, coltype in GET_COLUMNS_TYPE_MAP.items()
446+
]
447+
448+
SampleTable = Table(table_name, metadata_obj, *args)
449+
450+
metadata_obj.create_all(db_engine)
451+
452+
yield table_name
453+
454+
metadata_obj.drop_all(db_engine)
455+
456+
457+
def test_get_columns(db_engine, sample_table: str):
458+
"""Created after PECO-1297 and Github Issue #295 to verify that get_columsn behaves like it should for all known SQLAlchemy types"""
459+
460+
inspector = Inspector.from_engine(db_engine)
461+
462+
# this raises an exception if `parse_column_info_from_tgetcolumnsresponse` fails a lookup
463+
columns = inspector.get_columns(sample_table)
464+
465+
assert True

0 commit comments

Comments
 (0)