From a4072a491e015c960ad86a820595a62ac7747e7a Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 30 Nov 2023 15:12:54 -0500 Subject: [PATCH 1/4] Add e2e test that fails with: _raw_col_type = re.search(pat, thrift_resp_row.TYPE_NAME).group(0).lower() # type: ignore > _col_type = GET_COLUMNS_TYPE_MAP[_raw_col_type] E KeyError: 'timestamp_ntz' Signed-off-by: Jesse Whitehouse --- .../sqlalchemy/test_local/e2e/test_basic.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/src/databricks/sqlalchemy/test_local/e2e/test_basic.py b/src/databricks/sqlalchemy/test_local/e2e/test_basic.py index 7d3cad511..5396326af 100644 --- a/src/databricks/sqlalchemy/test_local/e2e/test_basic.py +++ b/src/databricks/sqlalchemy/test_local/e2e/test_basic.py @@ -431,3 +431,35 @@ 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): + """Creates a sample table and returns its table name""" + + table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")) + + SampleTable = Table( + table_name, + metadata_obj, + Column("string_example", String(255)), + Column("integer_example", Integer), + Column("boolean_example", BOOLEAN), + Column("decimal_example", DECIMAL(10, 2)), + Column("date_example", Date), + Column("datetime_example", DateTime) + ) + + metadata_obj.create_all(db_engine) + + yield table_name + + metadata_obj.drop_all(db_engine) + +def test_get_columns(db_engine, metadata_obj: MetaData, sample_table: str): + """Confirms that we get back the same time we declared in a model and inserted using Core""" + + + from sqlalchemy.engine.reflection import Inspector + + inspector = Inspector.from_engine(db_engine) + + columns = inspector.get_columns(sample_table) From 3ef76e6dff7239517af590acdcaa2a1f29e47a0d Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 30 Nov 2023 15:51:36 -0500 Subject: [PATCH 2/4] Update TGetColumnsResp parsing to handling TIMESTAMP_NTZ Closes issue #295 Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/_parse.py | 5 +- .../sqlalchemy/test_local/e2e/test_basic.py | 52 +++++++++---------- .../sqlalchemy/test_local/test_parsing.py | 1 + 3 files changed, 31 insertions(+), 27 deletions(-) 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 5396326af..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,22 +432,20 @@ 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): - """Creates a sample table and returns its table name""" + """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")) - SampleTable = Table( - table_name, - metadata_obj, - Column("string_example", String(255)), - Column("integer_example", Integer), - Column("boolean_example", BOOLEAN), - Column("decimal_example", DECIMAL(10, 2)), - Column("date_example", Date), - Column("datetime_example", DateTime) - ) + 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) @@ -454,12 +453,13 @@ def sample_table(metadata_obj: MetaData, db_engine: Engine): metadata_obj.drop_all(db_engine) -def test_get_columns(db_engine, metadata_obj: MetaData, sample_table: str): - """Confirms that we get back the same time we declared in a model and inserted using Core""" - - from sqlalchemy.engine.reflection import Inspector +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 diff --git a/src/databricks/sqlalchemy/test_local/test_parsing.py b/src/databricks/sqlalchemy/test_local/test_parsing.py index f17814f94..38b11673b 100644 --- a/src/databricks/sqlalchemy/test_local/test_parsing.py +++ b/src/databricks/sqlalchemy/test_local/test_parsing.py @@ -152,3 +152,4 @@ def test_build_pk_dict(): def test_filter_dict_by_value(match, output): result = match_dte_rows_by_value(FMT_SAMPLE_DT_OUTPUT, match) assert result == output + From 9dcc17c6258fc1bf7b37fd168152c6dd48f857f8 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 30 Nov 2023 16:20:41 -0500 Subject: [PATCH 3/4] Black the source code Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/test_local/test_parsing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/databricks/sqlalchemy/test_local/test_parsing.py b/src/databricks/sqlalchemy/test_local/test_parsing.py index 38b11673b..f17814f94 100644 --- a/src/databricks/sqlalchemy/test_local/test_parsing.py +++ b/src/databricks/sqlalchemy/test_local/test_parsing.py @@ -152,4 +152,3 @@ def test_build_pk_dict(): def test_filter_dict_by_value(match, output): result = match_dte_rows_by_value(FMT_SAMPLE_DT_OUTPUT, match) assert result == output - From 80bf945ca4bf3442e89eafdb4380a6a78913545b Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Thu, 30 Nov 2023 16:52:18 -0500 Subject: [PATCH 4/4] update changelog Signed-off-by: Jesse Whitehouse --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) 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