diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py index 8183d7ca30f..ca0db02f68a 100644 --- a/superset/connectors/sqla/models.py +++ b/superset/connectors/sqla/models.py @@ -156,9 +156,8 @@ class TableColumn(Model, BaseColumn): if not grain: raise NotImplementedError( f'No grain spec for {time_grain} for database {db.database_name}') - expr = db.db_engine_spec.get_time_expr( - self.expression or self.column_name, - pdf, time_grain, grain) + col = db.db_engine_spec.get_timestamp_column(self.expression, self.column_name) + expr = db.db_engine_spec.get_time_expr(col, pdf, time_grain, grain) sqla_col = literal_column(expr, type_=DateTime) return self.table.make_sqla_column_compatible(sqla_col, label) diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py index 07850820ad4..ad92fa9edc2 100644 --- a/superset/db_engine_specs.py +++ b/superset/db_engine_specs.py @@ -463,6 +463,13 @@ class BaseEngineSpec(object): label = label[:cls.max_column_name_length] return label + @staticmethod + def get_timestamp_column(expression, column_name): + """Return the expression if defined, otherwise return column_name. Some + engines require forcing quotes around column name, in which case this method + can be overridden.""" + return expression or column_name + class PostgresBaseEngineSpec(BaseEngineSpec): """ Abstract class for Postgres 'like' databases """ @@ -509,6 +516,16 @@ class PostgresEngineSpec(PostgresBaseEngineSpec): tables.extend(inspector.get_foreign_table_names(schema)) return sorted(tables) + @staticmethod + def get_timestamp_column(expression, column_name): + """Postgres is unable to identify mixed case column names unless they + are quoted.""" + if expression: + return expression + elif column_name.lower() != column_name: + return f'"{column_name}"' + return column_name + class SnowflakeEngineSpec(PostgresBaseEngineSpec): engine = 'snowflake' diff --git a/tests/model_tests.py b/tests/model_tests.py index cc2b30c485e..0fe03de932a 100644 --- a/tests/model_tests.py +++ b/tests/model_tests.py @@ -117,6 +117,39 @@ class DatabaseModelTestCase(SupersetTestCase): self.assertEquals(d.get('P1D').function, 'DATE({col})') self.assertEquals(d.get('Time Column').function, '{col}') + def test_postgres_expression_time_grain(self): + uri = 'postgresql+psycopg2://uid:pwd@localhost:5432/superset' + database = Database(sqlalchemy_uri=uri) + pdf, time_grain = '', 'P1D' + expression, column_name = 'COALESCE(lowercase_col, "MixedCaseCol")', '' + grain = database.grains_dict().get(time_grain) + col = database.db_engine_spec.get_timestamp_column(expression, column_name) + grain_expr = database.db_engine_spec.get_time_expr(col, pdf, time_grain, grain) + grain_expr_expected = grain.function.replace('{col}', expression) + self.assertEqual(grain_expr, grain_expr_expected) + + def test_postgres_lowercase_col_time_grain(self): + uri = 'postgresql+psycopg2://uid:pwd@localhost:5432/superset' + database = Database(sqlalchemy_uri=uri) + pdf, time_grain = '', 'P1D' + expression, column_name = '', 'lowercase_col' + grain = database.grains_dict().get(time_grain) + col = database.db_engine_spec.get_timestamp_column(expression, column_name) + grain_expr = database.db_engine_spec.get_time_expr(col, pdf, time_grain, grain) + grain_expr_expected = grain.function.replace('{col}', column_name) + self.assertEqual(grain_expr, grain_expr_expected) + + def test_postgres_mixedcase_col_time_grain(self): + uri = 'postgresql+psycopg2://uid:pwd@localhost:5432/superset' + database = Database(sqlalchemy_uri=uri) + pdf, time_grain = '', 'P1D' + expression, column_name = '', 'MixedCaseCol' + grain = database.grains_dict().get(time_grain) + col = database.db_engine_spec.get_timestamp_column(expression, column_name) + grain_expr = database.db_engine_spec.get_time_expr(col, pdf, time_grain, grain) + grain_expr_expected = grain.function.replace('{col}', f'"{column_name}"') + self.assertEqual(grain_expr, grain_expr_expected) + def test_single_statement(self): main_db = get_main_database(db.session)