Skip to content

Commit e0d54bb

Browse files
author
James Robinson
authored
Eng 5808 fix introspect databricks (#96)
1 parent 0fbeffe commit e0d54bb

File tree

1 file changed

+16
-10
lines changed

1 file changed

+16
-10
lines changed

noteable_magics/sql/meta_commands.py

Lines changed: 16 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from sqlalchemy import inspect
1818
from sqlalchemy.engine.reflection import Inspector
1919
from sqlalchemy.exc import NoSuchTableError
20+
from sqlalchemy.types import TypeEngine
2021

2122
from noteable_magics.sql.connection import Connection
2223
from noteable_magics.sql.gate_messaging_types import (
@@ -424,11 +425,7 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
424425
names.append(col['name'])
425426

426427
# Convert the possibly db-centric TypeEngine instance to a sqla-generic type string
427-
try:
428-
type_name = str(col['type'].as_generic()).lower()
429-
except NotImplementedError:
430-
# ENG-5268: More esoteric types like UUID do not implement .as_generic()
431-
type_name = str(col['type']).replace('()', '').lower()
428+
type_name = determine_column_type_name(col['type'])
432429

433430
types.append(type_name)
434431
nullables.append(col['nullable'])
@@ -805,11 +802,7 @@ def introspect_columns(
805802
for col in column_dicts:
806803
comment = col.get('comment') # Some dialects do not return.
807804

808-
try:
809-
type_name = str(col['type'].as_generic()).lower()
810-
except NotImplementedError:
811-
# ENG-5268: More esoteric types like UUID do not implement .as_generic()
812-
type_name = str(col['type']).replace('()', '').lower()
805+
type_name = determine_column_type_name(col['type'])
813806

814807
retlist.append(
815808
ColumnModel(
@@ -1297,6 +1290,19 @@ def _raise_from_no_such_table(schema: str, relation_name: str):
12971290
raise MetaCommandException(msg)
12981291

12991292

1293+
def determine_column_type_name(sqla_column_type_object: TypeEngine) -> str:
1294+
"""Convert the possibly db-centric TypeEngine instance to a sqla-generic type string"""
1295+
try:
1296+
type_name = str(sqla_column_type_object.as_generic()).lower()
1297+
except (NotImplementedError, AssertionError):
1298+
# ENG-5268: More esoteric types like UUID do not implement .as_generic()
1299+
# ENG-5808: Some Databricks types are not fully implemented and fail
1300+
# assertions within .as_generic()
1301+
type_name = str(sqla_column_type_object).replace('()', '').lower()
1302+
1303+
return type_name
1304+
1305+
13001306
def make_introspection_error_human_presentable(exception: Exception) -> str:
13011307
"""Convert any exception encountered by introspection into database into a nice human presentable string."""
13021308

0 commit comments

Comments
 (0)