Do only one query for all columns, instead of len(tables) queries. Should fix #25

pgcli/main.py

@@ -230,11 +230,11 @@ def quit_command(sql):
             or sql.strip() == ':q')
 
 def refresh_completions(pgexecute, completer):
-    tables = pgexecute.tables()
+    tables, columns = pgexecute.tables()
     completer.extend_table_names(tables)
     for table in tables:
         table = table[1:-1] if table[0] == '"' and table[-1] == '"' else table
-        completer.extend_column_names(table, pgexecute.columns(table))
+        completer.extend_column_names(table, columns[table])
     completer.extend_database_names(pgexecute.databases())
 
 if __name__ == "__main__":

pgcli/pgexecute.py

@@ -1,6 +1,7 @@
 import logging
 import psycopg2
 import psycopg2.extensions
+from collections import defaultdict
 from .packages import pgspecial
 
 _logger = logging.getLogger(__name__)
@@ -53,8 +54,7 @@ class PGExecute(object):
     'information_schema' AND n.nspname !~ '^pg_toast' AND
     pg_catalog.pg_table_is_visible(c.oid) ORDER BY 1;'''
 
-    columns_query = '''SELECT column_name FROM information_schema.columns WHERE
-    table_name =%s;'''
+    columns_query = '''SELECT table_name, column_name FROM information_schema.columns'''
 
     databases_query = """SELECT d.datname as "Name",
        pg_catalog.pg_get_userbyid(d.datdba) as "Owner",
@@ -123,21 +123,19 @@ class PGExecute(object):
                 return [(None, None, cur.statusmessage)]
 
     def tables(self):
+        """ Returns tuple (sorted_tables, columns). Columns is a dictionary of
+            table name -> list of columns """
+        columns = defaultdict(list)
         with self.conn.cursor() as cur:
             cur.execute(self.tables_query)
-            return [x[0] for x in cur.fetchall()]
+            tables = [x[0] for x in cur.fetchall()]
 
-    def columns(self, table):
-        with self.conn.cursor() as cur:
-            cur.execute(self.columns_query, (table,))
-            cols = [x[0] for x in cur.fetchall()]
-            return cols
-
-    def all_columns(self):
-        columns = set()
-        for table in self.tables():
-            columns.update(self.columns(table))
-        return columns
+            table_set = set(tables)
+            cur.execute(self.columns_query)
+            for table, column in cur.fetchall():
+                if table in table_set:
+                    columns[table].append(column)
+        return tables, columns
 
     def databases(self):
         with self.conn.cursor() as cur:

tests/test_pgexecute.py

@@ -51,3 +51,18 @@ def test_conn(executor):
         | abc |
         +-----+
         SELECT 1""")
+
+@dbtest
+def test_table_and_columns_query(executor):
+    run(executor, "create table a(x text, y text)")
+    run(executor, "create table b(z text)")
+
+    tables, columns = executor.tables()
+    assert tables == ['a', 'b']
+    assert columns['a'] == ['x', 'y']
+    assert columns['b'] == ['z']
+
+@dbtest
+def test_database_list(executor):
+    databases = executor.databases()
+    assert '_test_db' in databases