Merge pull request #880 from RishiRamraj/feature/get-last-sql-query

Feature/get last sql query

7 files changed, 64 insertions, 24 deletions

diff --git a/.gitignore b/.gitignore
index 75e1e8bf..55b4a74f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -37,9 +37,11 @@ pip-delete-this-directory.txt
 htmlcov/
 .tox/
 .coverage
+.coverage.*
 .cache
 nosetests.xml
 coverage.xml
+.pytest_cache
 
 # Translations
 *.mo
diff --git a/AUTHORS b/AUTHORS
index 74bf99d5..df481086 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -78,6 +78,7 @@ Contributors:
     * Dan Clark
     * Catherine Devlin
     * Jason Ribeiro
+    * Rishi Ramraj
 
 
 Creator:
diff --git a/changelog.rst b/changelog.rst
index 0be97c6a..d8c21080 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -11,6 +11,7 @@ Internal changes:
 
 * Mark tests requiring a running database server as dbtest (Thanks: `Dick Marinus`_)
 * Add ``application_name`` to help identify pgcli connection to database (issue #868) (Thanks: `François Pietka`_)
+* Add an is_special command flag to MetaQuery (Thanks: `Rishi Ramraj`_)
 * Ported Destructive Warning from mycli.
 
 Bug Fixes:
@@ -826,3 +827,4 @@ Improvements:
 .. _`Bojan Delić`: https://github.com/delicb
 .. _`Frederic Aoustin`: https://github.com/fraoustin
 .. _`Jason Ribeiro`: https://github.com/jrib
+.. _`Rishi Ramraj`: https://github.com/RishiRamraj
diff --git a/pgcli/main.py b/pgcli/main.py
index 1ce3f4b0..8ae5927d 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -80,6 +80,7 @@ MetaQuery = namedtuple(
         'db_changed',       # True if any subquery changed the database
         'path_changed',     # True if any subquery changed the search path
         'mutated',          # True if any subquery executed insert/update/delete
+        'is_special',       # True if the query is a special command
     ])
 MetaQuery.__new__.__defaults__ = ('', False, 0, False, False, False, False)
 
@@ -290,12 +291,12 @@ class PGCli(object):
     def execute_from_file(self, pattern, **_):
         if not pattern:
             message = '\\i: missing required argument'
-            return [(None, None, None, message, '', False)]
+            return [(None, None, None, message, '', False, True)]
         try:
             with open(os.path.expanduser(pattern), encoding='utf-8') as f:
                 query = f.read()
         except IOError as e:
-            return [(None, None, None, str(e), '', False)]
+            return [(None, None, None, str(e), '', False, True)]
 
         if (self.destructive_warning and
                 confirm_destructive_query(query) is False):
@@ -311,7 +312,7 @@ class PGCli(object):
         if not pattern:
             self.output_file = None
             message = 'File output disabled'
-            return [(None, None, None, message, '', True)]
+            return [(None, None, None, message, '', True, True)]
         filename = os.path.abspath(os.path.expanduser(pattern))
         if not os.path.isfile(filename):
             try:
@@ -319,10 +320,10 @@ class PGCli(object):
             except IOError as e:
                 self.output_file = None
                 message = str(e) + '\nFile output disabled'
-                return [(None, None, None, message, '', False)]
+                return [(None, None, None, message, '', False, True)]
         self.output_file = filename
         message = 'Writing to file "%s"' % self.output_file
-        return [(None, None, None, message, '', True)]
+        return [(None, None, None, message, '', True, True)]
 
     def initialize_logging(self):
 
@@ -490,9 +491,11 @@ class PGCli(object):
             cli.application.pre_run_callables = saved_callables
         return document
 
-    def execute_command(self, text, query):
+    def execute_command(self, text):
         logger = self.logger
 
+        query = MetaQuery(query=text, successful=False)
+
         try:
             if (self.destructive_warning):
                 destroy = confirm = confirm_destructive_query(text)
@@ -591,16 +594,12 @@ class PGCli(object):
                     click.secho(str(e), err=True, fg='red')
                     continue
 
-                # Initialize default metaquery in case execution fails
-                query = MetaQuery(query=document.text, successful=False)
-
                 self.watch_command, timing = special.get_watch_command(
                         document.text)
                 if self.watch_command:
                     while self.watch_command:
                         try:
-                            query = self.execute_command(
-                                    self.watch_command, query)
+                            query = self.execute_command(self.watch_command)
                             click.echo(
                                     'Waiting for {0} seconds before repeating'
                                     .format(timing))
@@ -608,7 +607,7 @@ class PGCli(object):
                         except KeyboardInterrupt:
                             self.watch_command = None
                 else:
-                    query = self.execute_command(document.text, query)
+                    query = self.execute_command(document.text)
 
                 self.now = dt.datetime.today()
 
@@ -725,7 +724,7 @@ class PGCli(object):
         res = self.pgexecute.run(text, self.pgspecial,
                                  exception_formatter, on_error_resume)
 
-        for title, cur, headers, status, sql, success in res:
+        for title, cur, headers, status, sql, success, is_special in res:
             logger.debug("headers: %r", headers)
             logger.debug("rows: %r", cur)
             logger.debug("status: %r", status)
@@ -772,7 +771,7 @@ class PGCli(object):
                 all_success = False
 
         meta_query = MetaQuery(text, all_success, total, meta_changed,
-                               db_changed, path_changed, mutated)
+                               db_changed, path_changed, mutated, is_special)
 
         return output, meta_query
 
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 2158f909..0d6195e0 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -286,7 +286,7 @@ class PGExecute(object):
                execute.
 
         :return: Generator yielding tuples containing
-                 (title, rows, headers, status, query, success)
+                 (title, rows, headers, status, query, success, is_special)
         """
 
         # Remove spaces and EOL
@@ -307,8 +307,8 @@ class PGExecute(object):
                     try:
                         for result in pgspecial.execute(cur, sql):
                             # e.g. execute_from_file already appends these
-                            if len(result) < 6:
-                                yield result + (sql, True)
+                            if len(result) < 7:
+                                yield result + (sql, True, True)
                             else:
                                 yield result
                         continue
@@ -316,7 +316,7 @@ class PGExecute(object):
                         pass
 
                 # Not a special command, so execute as normal sql
-                yield self.execute_normal_sql(sql) + (sql, True)
+                yield self.execute_normal_sql(sql) + (sql, True, False)
             except psycopg2.DatabaseError as e:
                 _logger.error("sql: %r, error: %r", sql, e)
                 _logger.error("traceback: %r", traceback.format_exc())
@@ -325,7 +325,7 @@ class PGExecute(object):
                         or not exception_formatter):
                     raise
 
-                yield None, None, None, exception_formatter(e), sql, False
+                yield None, None, None, exception_formatter(e), sql, False, False
 
                 if not on_error_resume:
                     break
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index 3141dc0e..56067906 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -2,10 +2,11 @@
 
 import pytest
 import psycopg2
-from pgspecial.main import PGSpecial
+from mock import patch
 from pgcli.packages.parseutils.meta import FunctionMetadata
 from textwrap import dedent
 from utils import run, dbtest, requires_json, requires_jsonb
+from pgcli.main import PGCli
 
 
 def function_meta_data(
@@ -179,6 +180,41 @@ def test_unicode_support_in_output(executor, expanded):
 
 
 @dbtest
+def test_not_is_special(executor, pgspecial):
+    """is_special is set to false for database queries."""
+    query = 'select 1'
+    result = list(executor.run(query, pgspecial=pgspecial))
+    success, is_special = result[0][5:]
+    assert success == True
+    assert is_special == False
+
+
+@dbtest
+def test_execute_from_file_no_arg(executor, pgspecial):
+    """\i without a filename returns an error."""
+    result = list(executor.run("\i", pgspecial=pgspecial))
+    status, sql, success, is_special = result[0][3:]
+    assert 'missing required argument' in status
+    assert success == False
+    assert is_special == True
+
+
+@dbtest
+@patch('pgcli.main.os')
+def test_execute_from_file_io_error(os, executor, pgspecial):
+    """\i with an io_error returns an error."""
+    # Inject an IOError.
+    os.path.expanduser.side_effect = IOError('test')
+
+    # Check the result.
+    result = list(executor.run("\i test", pgspecial=pgspecial))
+    status, sql, success, is_special = result[0][3:]
+    assert status == 'test'
+    assert success == False
+    assert is_special == True
+
+
+@dbtest
 def test_multiple_queries_same_line(executor):
     result = run(executor, "select 'foo'; select 'bar'")
     assert len(result) == 12  # 2 * (output+status) * 3 lines
@@ -205,7 +241,7 @@ def test_multiple_queries_same_line_syntaxerror(executor, exception_formatter):
 
 @pytest.fixture
 def pgspecial():
-    return PGSpecial()
+    return PGCli().pgspecial
 
 
 @dbtest
@@ -287,11 +323,11 @@ def test_large_numbers_render_directly(executor, value):
 @pytest.mark.parametrize('command', ['di', 'dv', 'ds', 'df', 'dT'])
 @pytest.mark.parametrize('verbose', ['', '+'])
 @pytest.mark.parametrize('pattern', ['', 'x', '*.*', 'x.y', 'x.*', '*.y'])
-def test_describe_special(executor, command, verbose, pattern):
+def test_describe_special(executor, command, verbose, pattern, pgspecial):
     # We don't have any tests for the output of any of the special commands,
     # but we can at least make sure they run without error
     sql = r'\{command}{verbose} {pattern}'.format(**locals())
-    executor.run(sql)
+    list(executor.run(sql, pgspecial=pgspecial))
 
 
 @dbtest
diff --git a/tests/utils.py b/tests/utils.py
index c8cb42bd..31fd12a3 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -68,7 +68,7 @@ def run(executor, sql, join=False, expanded=False, pgspecial=None,
     formatted = []
     settings = OutputSettings(table_format='psql', dcmlfmt='d', floatfmt='g',
                               expanded=expanded)
-    for title, rows, headers, status, sql, success in results:
+    for title, rows, headers, status, sql, success, is_special in results:
         formatted.extend(format_output(title, rows, headers, status, settings))
     if join:
         formatted = '\n'.join(formatted)