Postgres: make sure table from the public schema doesn't get merged with table from other schemas

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):
         '''

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'])