Move refresh routines outside the loop.

2 files changed, 25 insertions, 35 deletions

diff --git a/TODO b/TODO
index 00ff5130..21735457 100644
--- a/TODO
+++ b/TODO
@@ -7,12 +7,6 @@
 * [ ] Check how to add the name of the table before printing the table.
 * [ ] Add a new trigger for M-/ that does naive completion.
 * [ ] New Feature List - Write the current version to config file. At launch if the version has changed, display the changelog between the two versions.
-* [ ] need_search_path_refresh doesn't need to split the document.text, cur.query returns the current sql. 
-* [ ] See if the need_search_path_refresh can be done outside the loop.
 * [ ] Add a test for 'select * from custom.abc where custom.abc.' should suggest columns from abc.
-* [ ] Search for Dataframe in the source and replace them.
 * [ ] pgexecute columns(), tables() etc can be just cursors instead of fetchall()
-* [ ] Add unicode tests.
 * [ ] Add colorschemes in config file.
-* [X] Add \timing in config file.
-* [X] Add table format to config file.
diff --git a/pgcli/main.py b/pgcli/main.py
index 093bfca0..fa65e7f2 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -5,8 +5,10 @@ from __future__ import print_function
 import os
 import traceback
 import logging
-import click
+from time import time
 
+import click
+import sqlparse
 from prompt_toolkit import CommandLineInterface, AbortAction, Exit
 from prompt_toolkit.document import Document
 from prompt_toolkit.layout import Layout
@@ -29,7 +31,6 @@ from .config import write_default_config, load_config
 from .key_bindings import pgcli_bindings
 from .encodingutils import utf8tounicode
 
-from time import time
 
 try:
     from urlparse import urlparse
@@ -193,11 +194,12 @@ class PGCli(object):
                 try:
                     logger.debug('sql: %r', document.text)
                     successful = False
-                    start = time()
                     # Initialized to [] because res might never get initialized
                     # if an exception occurs in pgexecute.run(). Which causes
                     # finally clause to fail.
                     res = []
+                    start = time()
+                    # Run the query.
                     res = pgexecute.run(document.text)
                     duration = time() - start
                     successful = True
@@ -222,11 +224,6 @@ class PGCli(object):
                         total += end - start
                         mutating = mutating or is_mutating(status)
 
-                        if need_search_path_refresh(document.text, status):
-                            logger.debug('Refreshing search path')
-                            completer.set_search_path(pgexecute.search_path())
-                            logger.debug('Search path: %r', completer.search_path)
-
                 except KeyboardInterrupt:
                     # Restart connection to the database
                     pgexecute.connect()
@@ -251,6 +248,12 @@ class PGCli(object):
                     prompt = '%s> ' % pgexecute.dbname
                     self.refresh_completions()
 
+                # Refresh search_path to set default schema.
+                if need_search_path_refresh(document.text):
+                    logger.debug('Refreshing search path')
+                    completer.set_search_path(pgexecute.search_path())
+                    logger.debug('Search path: %r', completer.search_path)
+
                 query = Query(document.text, successful, mutating)
                 self.query_history.append(query)
 
@@ -337,28 +340,21 @@ def format_output(cur, headers, status, table_format):
         output.append(status)
     return output
 
-def need_completion_refresh(sql):
-    try:
-        first_token = sql.split()[0]
-        return first_token.lower() in ('alter', 'create', 'use', '\c', 'drop')
-    except Exception:
-        return False
-
-def need_search_path_refresh(sql, status):
-    # note that sql may be a multi-command query, but status belongs to an
-    # individual query, since pgexecute handles splitting up multi-commands
-    try:
-        status = status.split()[0]
-        if status.lower() == 'set':
-            # Since sql could be a multi-line query, it's hard to robustly
-            # pick out the variable name that's been set. Err on the side of
-            # false positives here, since the worst case is we refresh the
-            # search path when it's not necessary
-            return 'search_path' in sql.lower()
-        else:
+def need_completion_refresh(queries):
+    """Determines if the completion needs a refresh by checking if the sql
+    statement is an alter, create, drop or change db."""
+    for query in sqlparse.split(queries):
+        try:
+            first_token = query.split()[0]
+            return first_token.lower() in ('alter', 'create', 'use', '\c',
+                    'drop')
+        except Exception:
             return False
-    except Exception:
-        return False
+
+def need_search_path_refresh(sql):
+    """Determines if the search_path should be refreshed by checking if the
+    sql has 'set search_path'."""
+    return 'set search_path' in sql.lower()
 
 def is_mutating(status):
     """Determines if the statement is mutating based on the status."""