Issue 1018 display first 1k rows (#1092)

* Added changes to remove the prompt on >1000 rows queries

* Reformatted with black

* Changed comment on row_limit parameter

* Added contribution to changelog and name to AUTHORS

* Refactored test to reflect new functionality

* Removed argument

* Removed debug echo statement

* Reformatted with black

* Added changes to remove the prompt on >1000 rows queries

* Reformatted with black

* Changed comment on row_limit parameter

* Added contribution to changelog and name to AUTHORS

* Refactored test to reflect new functionality

* Removed argument

* Removed debug echo statement

* Reformatted with black

* Added missing issue numbers in changelog

* Reformatted code using black

4 files changed, 48 insertions, 38 deletions

diff --git a/changelog.rst b/changelog.rst
index 571f88cd..dbfa2749 100644
--- a/changelog.rst
+++ b/changelog.rst
@@ -5,12 +5,13 @@ Features:
 ---------
 
 * Add `\\G` as a terminator to sql statements that will show the results in expanded mode. This feature is copied from mycli. (Thanks: `Amjith Ramanujam`_)
+* Removed limit prompt and added automatic row limit on queries with no LIMIT clause (#1079) (Thanks: `Sebastian Janko`_)
 
 Bug fixes:
 ----------
 
 * Error connecting to PostgreSQL 12beta1 (#1058). (Thanks: `Irina Truong`_)
-* Empty query caused error message (Thanks: `Sebastian Janko`_)
+* Empty query caused error message (#1019) (Thanks: `Sebastian Janko`_)
 * History navigation bindings in multiline queries (#1004) (Thanks: `Pedro Ferrari`_)
 
 2.1.1
diff --git a/pgcli/main.py b/pgcli/main.py
index 27ca09b7..903966f4 100644
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -124,7 +124,6 @@ class PgCliQuitError(Exception):
 
 
 class PGCli(object):
-
     default_prompt = "\\u@\\h:\\d> "
     max_len_prompt = 30
 
@@ -828,11 +827,30 @@ class PGCli(object):
 
             return prompt_app
 
-    def _should_show_limit_prompt(self, status, cur):
-        """returns True if limit prompt should be shown, False otherwise."""
-        if not is_select(status):
+    def _should_limit_output(self, sql, cur):
+        """returns True if the output should be truncated, False otherwise."""
+        if not is_select(sql):
+            return False
+
+        return (
+            not self._has_limit(sql)
+            and self.row_limit != 0
+            and cur
+            and cur.rowcount > self.row_limit
+        )
+
+    def _has_limit(self, sql):
+        if not sql:
             return False
-        return self.row_limit > 0 and cur and (cur.rowcount > self.row_limit)
+        return "limit " in sql.lower()
+
+    def _limit_output(self, cur):
+        limit = min(self.row_limit, cur.rowcount)
+        new_cur = itertools.islice(cur, limit)
+        new_status = "SELECT " + str(limit)
+        click.secho("The result was limited to %s rows" % limit, fg="red")
+
+        return new_cur, new_status
 
     def _evaluate_command(self, text):
         """Used to run a command entered by the user during CLI operation
@@ -865,14 +883,9 @@ class PGCli(object):
             logger.debug("headers: %r", headers)
             logger.debug("rows: %r", cur)
             logger.debug("status: %r", status)
-            threshold = self.row_limit
-            if self._should_show_limit_prompt(status, cur):
-                click.secho(
-                    "The result set has more than %s rows." % threshold, fg="red"
-                )
-                if not click.confirm("Do you want to continue?"):
-                    click.secho("Aborted!", err=True, fg="red")
-                    break
+
+            if self._should_limit_output(sql, cur):
+                cur, status = self._limit_output(cur)
 
             if self.pgspecial.auto_expand or self.auto_expand:
                 max_width = self.prompt_app.output.get_size().columns
@@ -1184,7 +1197,6 @@ def cli(
     list_dsn,
     warn,
 ):
-
     if version:
         print ("Version:", __version__)
         sys.exit(0)
@@ -1416,12 +1428,12 @@ def format_output(title, cur, headers, status, settings):
                     column_types.append(int)
                 else:
                     column_types.append(text_type)
+
         formatted = formatter.format_output(cur, headers, **output_kwargs)
         if isinstance(formatted, (text_type)):
             formatted = iter(formatted.splitlines())
         first_line = next(formatted)
         formatted = itertools.chain([first_line], formatted)
-
         if not expanded and max_width and len(first_line) > max_width and headers:
             formatted = formatter.format_output(
                 cur, headers, format_name="vertical", column_types=None, **output_kwargs
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index 78d78366..29b1b4c5 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -109,7 +109,7 @@ vi = False
 # Possible values "STOP" or "RESUME"
 on_error = STOP
 
-# Set threshold for row limit prompt. Use 0 to disable prompt.
+# Set threshold for row limit. Use 0 to disable limiting.
 row_limit = 1000
 
 # Skip intro on startup and goodbye on exit
diff --git a/tests/test_rowlimit.py b/tests/test_rowlimit.py
index bc66f1e7..e76ea049 100644
--- a/tests/test_rowlimit.py
+++ b/tests/test_rowlimit.py
@@ -1,6 +1,7 @@
-from pgcli.main import PGCli
-from mock import Mock
 import pytest
+from mock import Mock
+
+from pgcli.main import PGCli
 
 
 # We need this fixtures beacause we need PGCli object to be created
@@ -43,40 +44,36 @@ def low_count():
     return low_count_cursor
 
 
-def test_default_row_limit(low_count, over_default):
-    cli = PGCli()
-    stmt = "SELECT * FROM students"
-    result = cli._should_show_limit_prompt(stmt, low_count)
+def test_row_limit_with_LIMIT_clause(LIMIT, over_limit):
+    cli = PGCli(row_limit=LIMIT)
+    stmt = "SELECT * FROM students LIMIT 1000"
+
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False
 
-    result = cli._should_show_limit_prompt(stmt, over_default)
-    assert result is True
+    cli = PGCli(row_limit=0)
+    result = cli._should_limit_output(stmt, over_limit)
+    assert result is False
 
 
-def test_set_row_limit(over_default, over_limit, LIMIT):
+def test_row_limit_without_LIMIT_clause(LIMIT, over_limit):
     cli = PGCli(row_limit=LIMIT)
     stmt = "SELECT * FROM students"
-    result = cli._should_show_limit_prompt(stmt, over_default)
-    assert result is False
 
-    result = cli._should_show_limit_prompt(stmt, over_limit)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is True
 
-
-def test_no_limit(over_limit):
     cli = PGCli(row_limit=0)
-    stmt = "SELECT * FROM students"
-
-    result = cli._should_show_limit_prompt(stmt, over_limit)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False
 
 
-def test_row_limit_on_non_select(over_default):
+def test_row_limit_on_non_select(over_limit):
     cli = PGCli()
-    stmt = "UPDATE students set name='Boby'"
-    result = cli._should_show_limit_prompt(stmt, None)
+    stmt = "UPDATE students SET name='Boby'"
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False
 
     cli = PGCli(row_limit=0)
-    result = cli._should_show_limit_prompt(stmt, over_default)
+    result = cli._should_limit_output(stmt, over_limit)
     assert result is False