Skip to content
This repository has been archived by the owner on Nov 3, 2023. It is now read-only.

Commit

Permalink
Force lowercase column names for Snowflake and Oracle (apache#4994)
Browse files Browse the repository at this point in the history
* Force lowercase column names for Snowflake and Oracle

* Force lowercase column names for Snowflake and Oracle

* Remove lowercasing of DB2 columns

* Remove DB2 lowercasing

* Fix test cases
  • Loading branch information
villebro authored and timifasubaa committed Jul 25, 2018
1 parent d59e564 commit 4502ec7
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 8 deletions.
17 changes: 17 additions & 0 deletions superset/db_engine_specs.py
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,15 @@ def get_configuration_for_impersonation(cls, uri, impersonate_user, username):
"""
return {}

@classmethod
def get_normalized_column_names(cls, cursor_description):
columns = cursor_description if cursor_description else []
return [cls.normalize_column_name(col[0]) for col in columns]

@staticmethod
def normalize_column_name(column_name):
return column_name


class PostgresBaseEngineSpec(BaseEngineSpec):
""" Abstract class for Postgres 'like' databases """
Expand Down Expand Up @@ -350,6 +359,10 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
Grain('year', _('year'), "DATE_TRUNC('YEAR', {col})", 'P1Y'),
)

@staticmethod
def normalize_column_name(column_name):
return column_name.lower()


class VerticaEngineSpec(PostgresBaseEngineSpec):
engine = 'vertica'
Expand Down Expand Up @@ -379,6 +392,10 @@ def convert_dttm(cls, target_type, dttm):
"""TO_TIMESTAMP('{}', 'YYYY-MM-DD"T"HH24:MI:SS.ff6')"""
).format(dttm.isoformat())

@staticmethod
def normalize_column_name(column_name):
return column_name.lower()


class Db2EngineSpec(BaseEngineSpec):
engine = 'ibm_db_sa'
Expand Down
8 changes: 3 additions & 5 deletions superset/sql_lab.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,10 +97,8 @@ def session_scope(nullpool):
session.close()


def convert_results_to_df(cursor_description, data):
def convert_results_to_df(column_names, data):
"""Convert raw query results to a DataFrame."""
column_names = (
[col[0] for col in cursor_description] if cursor_description else [])
column_names = dedup(column_names)

# check whether the result set has any nested dict columns
Expand Down Expand Up @@ -236,7 +234,7 @@ def handle_error(msg):
return handle_error(db_engine_spec.extract_error_message(e))

logging.info('Fetching cursor description')
cursor_description = cursor.description
column_names = db_engine_spec.get_normalized_column_names(cursor.description)

if conn is not None:
conn.commit()
Expand All @@ -245,7 +243,7 @@ def handle_error(msg):
if query.status == utils.QueryStatus.STOPPED:
return handle_error('The query has been stopped')

cdf = convert_results_to_df(cursor_description, data)
cdf = convert_results_to_df(column_names, data)

query.rows = cdf.size
query.progress = 100
Expand Down
6 changes: 3 additions & 3 deletions tests/sqllab_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,23 +203,23 @@ def test_alias_duplicate(self):
raise_on_error=True)

def test_df_conversion_no_dict(self):
cols = [['string_col'], ['int_col'], ['float_col']]
cols = ['string_col', 'int_col', 'float_col']
data = [['a', 4, 4.0]]
cdf = convert_results_to_df(cols, data)

self.assertEquals(len(data), cdf.size)
self.assertEquals(len(cols), len(cdf.columns))

def test_df_conversion_tuple(self):
cols = [['string_col'], ['int_col'], ['list_col'], ['float_col']]
cols = ['string_col', 'int_col', 'list_col', 'float_col']
data = [(u'Text', 111, [123], 1.0)]
cdf = convert_results_to_df(cols, data)

self.assertEquals(len(data), cdf.size)
self.assertEquals(len(cols), len(cdf.columns))

def test_df_conversion_dict(self):
cols = [['string_col'], ['dict_col'], ['int_col']]
cols = ['string_col', 'dict_col', 'int_col']
data = [['a', {'c1': 1, 'c2': 2, 'c3': 3}, 4]]
cdf = convert_results_to_df(cols, data)

Expand Down

0 comments on commit 4502ec7

Please sign in to comment.