diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py index e9e4cc5431..f6060d431d 100644 --- a/redash/query_runner/pg.py +++ b/redash/query_runner/pg.py @@ -63,6 +63,38 @@ def _wait(conn, timeout=None): raise psycopg2.OperationalError("select.error received") +def full_table_name(schema, name): + if '.' in name: + name = u'"{}"'.format(name) + + return u'{}.{}'.format(schema, name) + + +def build_schema(query_result, schema): + # By default we omit the public schema name from the table name. But there are + # edge cases, where this might cause conflicts. For example: + # * We have a schema named "main" with table "users". + # * We have a table named "main.users" in the public schema. + # (while this feels unlikely, this actually happened) + # In this case if we omit the schema name for the public table, we will have + # a conflict. + table_names = set(map(lambda r: full_table_name(r['table_schema'], r['table_name']), query_result['rows'])) + + for row in query_result['rows']: + if row['table_schema'] != 'public': + table_name = full_table_name(row['table_schema'], row['table_name']) + else: + if row['table_name'] in table_names: + table_name = full_table_name(row['table_schema'], row['table_name']) + else: + table_name = row['table_name'] + + if table_name not in schema: + schema[table_name] = {'name': table_name, 'columns': []} + + schema[table_name]['columns'].append(row['column_name']) + + class PostgreSQL(BaseSQLQueryRunner): noop_query = "SELECT 1" @@ -112,17 +144,7 @@ def _get_definitions(self, schema, query): results = json_loads(results) - for row in results['rows']: - if row['table_schema'] != 'public': - table_name = u'{}.{}'.format(row['table_schema'], - row['table_name']) - else: - table_name = row['table_name'] - - if table_name not in schema: - schema[table_name] = {'name': table_name, 'columns': []} - - schema[table_name]['columns'].append(row['column_name']) + build_schema(results, schema) def _get_tables(self, schema): ''' diff --git a/tests/query_runner/test_pg.py b/tests/query_runner/test_pg.py new file mode 100644 index 0000000000..6b244cfede --- /dev/null +++ b/tests/query_runner/test_pg.py @@ -0,0 +1,23 @@ + +from unittest import TestCase +from redash.query_runner.pg import build_schema + + +class TestBuildSchema(TestCase): + def test_handles_dups_between_public_and_other_schemas(self): + results = { + 'rows': [ + {'table_schema': 'public', 'table_name': 'main.users', 'column_name': 'id'}, + {'table_schema': 'main', 'table_name': 'users', 'column_name': 'id'}, + {'table_schema': 'main', 'table_name': 'users', 'column_name': 'name'}, + ] + } + + schema = {} + + build_schema(results, schema) + + self.assertIn('main.users', schema.keys()) + self.assertListEqual(schema['main.users']['columns'], ['id', 'name']) + self.assertIn('public."main.users"', schema.keys()) + self.assertListEqual(schema['public."main.users"']['columns'], ['id']) \ No newline at end of file