Merge pull request #521 from Smotko/configure-row-limit

Make too many rows warning limit configurable

3 files changed, 75 insertions, 7 deletions

diff --git a/pgcli/main.py b/pgcli/main.py
index c387e9c6..65e27018 100755
--- a/pgcli/main.py
+++ b/pgcli/main.py
@@ -85,7 +85,7 @@ class PGCli(object):
         if configured_pager:
             self.logger.info('Default pager found in config file: ' + '\'' + configured_pager + '\'')
             os.environ['PAGER'] = configured_pager
-        elif (os_environ_pager):
+        elif os_environ_pager:
             self.logger.info('Default pager found in PAGER environment variable: ' + '\'' + os_environ_pager + '\'')
             os.environ['PAGER'] = os_environ_pager
         else:
@@ -95,7 +95,7 @@ class PGCli(object):
         os.environ['LESS'] = '-SRXF'
 
     def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False,
-                 pgexecute=None, pgclirc_file=None):
+                 pgexecute=None, pgclirc_file=None, row_limit=None):
 
         self.force_passwd_prompt = force_passwd_prompt
         self.never_passwd_prompt = never_passwd_prompt
@@ -122,6 +122,10 @@ class PGCli(object):
         self.multi_line = c['main'].as_bool('multi_line')
         self.vi_mode = c['main'].as_bool('vi')
         self.pgspecial.timing_enabled = c['main'].as_bool('timing')
+        if row_limit is not None:
+            self.row_limit = row_limit
+        else:
+            self.row_limit = c['main'].as_int('row_limit')
 
         self.table_format = c['main']['table_format']
         self.syntax_style = c['main']['syntax_style']
@@ -492,6 +496,12 @@ class PGCli(object):
 
             return cli
 
+    def _should_show_limit_prompt(self, status, cur):
+        """returns True if limit prompt should be shown, False otherwise."""
+        if not is_select(status):
+            return False
+        return self.row_limit > 0 and cur and cur.rowcount > self.row_limit
+
     def _evaluate_command(self, text):
         """Used to run a command entered by the user during CLI operation
         (Puts the E in REPL)
@@ -519,9 +529,8 @@ class PGCli(object):
             logger.debug("headers: %r", headers)
             logger.debug("rows: %r", cur)
             logger.debug("status: %r", status)
-            threshold = 1000
-            if (is_select(status) and
-                    cur and cur.rowcount > threshold):
+            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?'):
@@ -655,10 +664,12 @@ class PGCli(object):
         envvar='PGCLIRC', help='Location of pgclirc file.')
 @click.option('-D', '--dsn', default='', envvar='DSN',
         help='Use DSN configured into the [alias_dsn] section of pgclirc file.')
+@click.option('-R', '--row-limit', default=None, envvar='PGROWLIMIT', type=click.INT,
+        help='Set threshold for row limit prompt. Use 0 to disable prompt.')
 @click.argument('database', default=lambda: None, envvar='PGDATABASE', nargs=1)
 @click.argument('username', default=lambda: None, envvar='PGUSER', nargs=1)
 def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
-        username, version, pgclirc, dsn):
+        username, version, pgclirc, dsn, row_limit):
 
     if version:
         print('Version:', __version__)
@@ -680,7 +691,8 @@ def cli(database, user, host, port, prompt_passwd, never_prompt, dbname,
             print ('Please move the existing config file ~/.pgclirc to',
                    config_full_path)
 
-    pgcli = PGCli(prompt_passwd, never_prompt, pgclirc_file=pgclirc)
+    pgcli = PGCli(prompt_passwd, never_prompt, pgclirc_file=pgclirc,
+                  row_limit=row_limit)
 
     # Choose which ever one has a valid value.
     database = database or dbname
diff --git a/pgcli/pgclirc b/pgcli/pgclirc
index da8c238c..ae7d6b09 100644
--- a/pgcli/pgclirc
+++ b/pgcli/pgclirc
@@ -64,6 +64,8 @@ vi = False
 # Possible values "STOP" or "RESUME"
 on_error = STOP
 
+# Set threshold for row limit prompt. Use 0 to disable prompt.
+row_limit = 1000
 
 # Custom colors for the completion menu, toolbar, etc.
 [colors]
diff --git a/tests/test_rowlimit.py b/tests/test_rowlimit.py
new file mode 100644
index 00000000..8d9a87a5
--- /dev/null
+++ b/tests/test_rowlimit.py
@@ -0,0 +1,54 @@
+from pgcli.main import PGCli
+from mock import Mock
+
+DEFAULT = PGCli().row_limit
+LIMIT = DEFAULT + 1000
+
+
+over_default = Mock()
+over_default.configure_mock(rowcount=DEFAULT + 10)
+
+over_limit = Mock()
+over_limit.configure_mock(rowcount=LIMIT + 10)
+
+low_count = Mock()
+low_count.configure_mock(rowcount=1)
+
+
+def test_default_row_limit():
+    cli = PGCli()
+    stmt = "SELECT * FROM students"
+    result = cli._should_show_limit_prompt(stmt, low_count)
+    assert result is False
+
+    result = cli._should_show_limit_prompt(stmt, over_default)
+    assert result is True
+
+
+def test_set_row_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)
+    assert result is True
+
+
+def test_no_limit():
+    cli = PGCli(row_limit=0)
+    stmt = "SELECT * FROM students"
+
+    result = cli._should_show_limit_prompt(stmt, over_limit)
+    assert result is False
+
+
+def test_row_limit_on_non_select():
+    cli = PGCli()
+    stmt = "UPDATE students set name='Boby'"
+    result = cli._should_show_limit_prompt(stmt, None)
+    assert result is False
+
+    cli = PGCli(row_limit=0)
+    result = cli._should_show_limit_prompt(stmt, over_default)
+    assert result is False