Cast row_limit as integer

mssqlcli/mssql_cli.py

@@ -167,11 +167,7 @@ def __init__(self, options):
             self.multiline_mode = c['main'].get('multi_line_mode', 'tsql')
             self.vi_mode = c['main'].as_bool('vi')
             self.prompt_format = options.prompt or c['main'].get('prompt', self.default_prompt)
-            if options.row_limit is not None:
-                self.row_limit = options.row_limit
-            else:
-                self.row_limit = c['main'].as_int('row_limit')
-
+            self.row_limit = options.row_limit
             self.min_num_menu_lines = c['main'].as_int('min_num_menu_lines')
             self.multiline_continuation_char = c['main']['multiline_continuation_char']
             self.syntax_style = c['main']['syntax_style']

mssqlcli/mssqlclioptionsparser.py

@@ -1,8 +1,10 @@
 import argparse
 import os
+import sys
+import click
 import mssqlcli
 
-from .config import config_location
+from .config import config_location, get_config
 
 MSSQL_CLI_USER = u'MSSQL_CLI_USER'
 MSSQL_CLI_PASSWORD = u'MSSQL_CLI_PASSWORD'
@@ -70,7 +72,9 @@ def create_parser():
     args_parser.add_argument(
         u'--row-limit',
         dest=u'row_limit',
-        default=os.environ.get(MSSQL_CLI_ROW_LIMIT, None),
+        default=get_config()['main'].as_int('row_limit') or \
+                os.environ.get(MSSQL_CLI_ROW_LIMIT, None),
+        type=check_row_limit,
         metavar=u'',
         help=u'Set threshold for row limit prompt. Use 0 to disable prompt.')
 
@@ -187,3 +191,23 @@ def create_parser():
     )
 
     return args_parser
+
+def check_row_limit(row_limit):
+    """
+    Validates row_limit option has valid positive integer
+    """
+
+    try:
+        row_limit_int = int(row_limit)
+    except ValueError:
+        click.secho(u'Error: row-limit has been set to an invalid value.\nPlease ' \
+                    u'specify a positive integer using the --row-limit command-line ' \
+                    u'argument or by setting row_limit in the config file.', fg='red')
+        sys.exit(1)
+    else:
+        if row_limit_int < 0:
+            click.secho(u'Error: row-limit cannot be negative.\nPlease ' \
+                        u'specify a positive integer using the --row-limit command-line ' \
+                        u'argument or by setting row_limit in the config file.', fg='red')
+            sys.exit(1)
+        return row_limit_int

tests/test_interactive_mode.py

@@ -27,6 +27,7 @@ def mssqlcli():
         yield mssqlcli
         shutdown(mssqlcli)
 
+
 class TestInteractiveModeQueries(TestInteractiveMode):
     @staticmethod
     @pytest.mark.parametrize("query_str, test_file", test_queries)
@@ -37,6 +38,7 @@ def test_query(query_str, test_file, mssqlcli):
         output_query = '\n'.join(mssqlcli.execute_query(query_str)).replace('\r', '')
         assert output_query == output_baseline
 
+
 class TestInteractiveModeInvalidRuns:
     @pytest.mark.timeout(60)
     def test_noninteractive_run(self):
@@ -83,6 +85,7 @@ def invalid_run(**options):
             if mssqlcli is not None:
                 shutdown(mssqlcli)
 
+
 class TestInteractiveModePager(TestInteractiveMode):
     """
     Test default pager setting.
@@ -120,6 +123,12 @@ def test_pager_config(mssqlcli):
         config['main']['pager'] = config_value
         assert mssqlcli.set_default_pager(config) == config_value
 
+
+class TestInteractiveModeRun:
+    """
+    Tests the executable.
+    """
+
     @staticmethod
     @pytest.mark.timeout(60)
     def test_valid_command():

tests/test_rowlimit.py

@@ -1,6 +1,8 @@
 # pylint: disable=protected-access
 import unittest
+import pytest
 from mssqlcli.mssql_cli import MssqlCli
+from mssqlcli.mssqlclioptionsparser import check_row_limit
 from mssqltestutils import create_mssql_cli_options
 
 class RowLimitTests(unittest.TestCase):
@@ -43,3 +45,35 @@ def test_row_limit_on_non_select(self):
         result = cli._should_show_limit_prompt(stmt, ['row']*self.over_default)
         assert cli.row_limit == 0
         assert not result
+
+
+class TestRowLimitArgs:
+    """
+    Tests for valid row-limit arguments.
+    """
+    # pylint: disable=protected-access
+
+    test_data_valid = [5, 0]
+    test_data_invalid = ['string!', -3]
+
+    @staticmethod
+    @pytest.mark.parametrize("row_limit", test_data_valid)
+    def test_valid_row_limit(row_limit):
+        """
+        Test valid value types for row limit argument
+        """
+        assert check_row_limit(row_limit) == row_limit
+
+    @staticmethod
+    @pytest.mark.parametrize("row_limit", test_data_invalid)
+    def test_invalid_row_limit(row_limit):
+        """
+        Test invalid value types for row limit argument
+        """
+        try:
+            check_row_limit(row_limit)
+        except SystemExit:
+            # mssqlcli class calls sys.exit(1) on invalid value
+            assert True
+        else:
+            assert False