Beispiel #1
0
class MssqlCli(object):

    max_len_prompt = 30
    default_prompt = '\\d> '

    def set_default_pager(self, config):
        configured_pager = config['main'].get('pager')
        os_environ_pager = os.environ.get('PAGER')

        if configured_pager:
            self.logger.info(
                'Default pager found in config file: "{}"'.format(configured_pager))
            os.environ['PAGER'] = configured_pager
        elif os_environ_pager:
            self.logger.info('Default pager found in PAGER environment variable: "{}"'.format(
                os_environ_pager))
            os.environ['PAGER'] = os_environ_pager
        else:
            self.logger.info(
                'No default pager found in environment. Using os default pager')

        # Set default set of less recommended options, if they are not already set.
        # They are ignored if pager is different than less.
        if not os.environ.get('LESS'):
            os.environ['LESS'] = '-SRXF'

    def __init__(self, options):

        # Load config.
        c = self.config = get_config(options.mssqlclirc_file)

        self.initialize_logging()
        self.logger = logging.getLogger(u'mssqlcli.main')

        self.set_default_pager(c)
        self.output_file = None

        self.multi_line = c['main'].as_bool('multi_line')
        self.multiline_mode = c['main'].get('multi_line_mode', 'tsql')
        self.vi_mode = c['main'].as_bool('vi')
        self.auto_expand = options.auto_vertical_output or c['main']['expand'] == 'auto'
        self.expanded_output = c['main']['expand'] == 'always'
        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.min_num_menu_lines = c['main'].as_int('min_num_menu_lines')
        self.multiline_continuation_char = c['main']['multiline_continuation_char']
        self.table_format = c['main']['table_format']
        self.syntax_style = c['main']['syntax_style']
        self.cli_style = c['colors']
        self.wider_completion_menu = c['main'].as_bool('wider_completion_menu')
        self.less_chatty = bool(
            options.less_chatty) or c['main'].as_bool('less_chatty')
        self.null_string = c['main'].get('null_string', '<null>')
        self.on_error = c['main']['on_error'].upper()
        self.decimal_format = c['data_formats']['decimal']
        self.float_format = c['data_formats']['float']

        self.now = dt.datetime.today()

        self.completion_refresher = CompletionRefresher()

        self.query_history = []

        # Initialize completer
        smart_completion = True if c['main'].get('smart_completion', 'True') == 'True' else False
        keyword_casing = c['main']['keyword_casing']
        self.settings = {
            'casing_file': get_casing_file(c),
            'generate_casing_file': c['main'].as_bool('generate_casing_file'),
            'generate_aliases': c['main'].as_bool('generate_aliases'),
            'asterisk_column_order': c['main']['asterisk_column_order'],
            'qualify_columns': c['main']['qualify_columns'],
            'case_column_headers': c['main'].as_bool('case_column_headers'),
            'search_path_filter': c['main'].as_bool('search_path_filter'),
            'single_connection': False,
            'less_chatty': options.less_chatty,
            'keyword_casing': keyword_casing,
        }

        self.completer = MssqlCompleter(smart_completion=smart_completion, settings=self.settings)
        self._completer_lock = threading.Lock()

        self.eventloop = create_eventloop()
        self.cli = None
        self.integrated_auth = options.integrated_auth

        self.sqltoolsclient = SqlToolsClient(enable_logging=options.enable_sqltoolsservice_logging)
        self.mssqlcliclient_main = MssqlCliClient(options, self.sqltoolsclient)

    def __del__(self):
        # Shut-down sqltoolsservice
        if self.sqltoolsclient:
            self.sqltoolsclient.shutdown()

    def write_to_file(self, pattern, **_):
        if not pattern:
            self.output_file = None
            message = 'File output disabled'
            return [(None, None, None, message, '', True)]
        filename = os.path.abspath(os.path.expanduser(pattern))
        if not os.path.isfile(filename):
            try:
                open(filename, 'w').close()
            except IOError as e:
                self.output_file = None
                message = str(e) + '\nFile output disabled'
                return [(None, None, None, message, '', False)]
        self.output_file = filename
        message = 'Writing to file "%s"' % self.output_file
        return [(None, None, None, message, '', True)]

    def initialize_logging(self):

        log_file = self.config['main']['log_file']
        if log_file == 'default':
            log_file = config_location() + 'mssqlcli.log'
        ensure_dir_exists(log_file)
        log_level = self.config['main']['log_level']

        # Disable logging if value is NONE by switching to a no-op handler.
        # Set log level to a high value so it doesn't even waste cycles getting
        # called.
        if log_level.upper() == 'NONE':
            handler = logging.NullHandler()
        else:
            handler = logging.FileHandler(os.path.expanduser(log_file))

        level_map = {'CRITICAL': logging.CRITICAL,
                     'ERROR': logging.ERROR,
                     'WARNING': logging.WARNING,
                     'INFO': logging.INFO,
                     'DEBUG': logging.DEBUG,
                     'NONE': logging.CRITICAL
                     }

        log_level = level_map[log_level.upper()]

        formatter = logging.Formatter(
            '%(asctime)s (%(process)d/%(threadName)s) '
            '%(name)s %(levelname)s - %(message)s')

        handler.setFormatter(formatter)

        root_logger = logging.getLogger('mssqlcli')
        root_logger.addHandler(handler)
        root_logger.setLevel(log_level)

        root_logger.info('Initializing mssqlcli logging.')
        root_logger.debug('Log file %r.', log_file)

    def set_main_mssqlcli_client(self, mssqlcli_client):
        self.mssqlcliclient_main = mssqlcli_client

    def connect_to_database(self):

        try:
            owner_uri, error_messages = self.mssqlcliclient_main.connect_to_database()
            if not owner_uri and error_messages:
                click.secho('\n'.join(error_messages),
                            err=True,
                            fg='yellow')
                sys.exit(1)

        except Exception as e:
            self.logger.debug('Database connection failed: %r.', e)
            self.logger.error("traceback: %r", traceback.format_exc())
            click.secho(str(e), err=True, fg='yellow')
            sys.exit(1)

    def handle_editor_command(self, cli, document):
        r"""
        Editor command is any query that is prefixed or suffixed
        by a '\e'. The reason for a while loop is because a user
        might edit a query multiple times.
        For eg:
        "select * from \e"<enter> to edit it in vim, then come
        back to the prompt with the edited query "select * from
        blah where q = 'abc'\e" to edit it again.
        :param cli: CommandLineInterface
        :param document: Document
        :return: Document
        """
        # FIXME: using application.pre_run_callables like this here is not the best solution.
        # It's internal api of prompt_toolkit that may change. This was added to fix #668.
        # We may find a better way to do it in the future.
        saved_callables = cli.application.pre_run_callables
        while special.editor_command(document.text):
            filename = special.get_filename(document.text)
            query = (special.get_editor_query(document.text) or
                     self.get_last_query())
            sql, message = special.open_external_editor(filename, sql=query)
            if message:
                # Something went wrong. Raise an exception and bail.
                raise RuntimeError(message)
            cli.current_buffer.document = Document(sql, cursor_position=len(sql))
            cli.application.pre_run_callables = []
            document = cli.run()
            continue
        cli.application.pre_run_callables = saved_callables
        return document

    def execute_command(self, text, query):
        logger = self.logger

        try:
            output, query = self._evaluate_command(text)
        except KeyboardInterrupt:
            # Issue where Ctrl+C propagates to sql tools service process and kills it,
            # so that query/cancel request can't be sent.
            # Right now the sql_tools_service process is killed and we restart
            # it with a new connection.
            click.secho(u'Cancelling query...', err=True, fg='red')
            self.reset()
            logger.debug("cancelled query, sql: %r", text)
            click.secho("Query cancelled.", err=True, fg='red')

        except NotImplementedError:
            click.secho('Not Yet Implemented.', fg="yellow")
        except Exception as e:
            logger.error("sql: %r, error: %r", text, e)
            logger.error("traceback: %r", traceback.format_exc())
            click.secho(str(e), err=True, fg='red')
        else:
            try:
                if self.output_file and not text.startswith(('\\o ', '\\? ')):
                    try:
                        with open(self.output_file, 'a', encoding='utf-8') as f:
                            click.echo(text, file=f)
                            click.echo('\n'.join(output), file=f)
                            click.echo('', file=f)  # extra newline
                    except IOError as e:
                        click.secho(str(e), err=True, fg='red')
                else:
                    click.echo_via_pager('\n'.join(output))
            except KeyboardInterrupt:
                pass

            if query.total_time > 1:
                print('Time: %0.03fs (%s)' % (query.total_time,
                                              humanize.time.naturaldelta(query.total_time)))
            else:
                print('Time: %0.03fs' % query.total_time)

            # Check if we need to update completions, in order of most
            # to least drastic changes
            if query.db_changed:
                with self._completer_lock:
                    self.completer.reset_completions()
                self.refresh_completions(persist_priorities='keywords')
            elif query.meta_changed:
                self.refresh_completions(persist_priorities='all')

        return query

    def run(self):
        history_file = self.config['main']['history_file']
        if history_file == 'default':
            history_file = config_location() + 'history'
        history = MssqlFileHistory(os.path.expanduser(history_file))

        self.refresh_completions(history=history,
                                 persist_priorities='none')

        self.cli = self._build_cli(history)

        if not self.less_chatty:
            print('Version: {}'.format(__version__))
            print('Mail: [email protected]')
            print('Home: http://github.com/dbcli/mssql-cli')

        try:
            while True:
                document = self.cli.run()

                # The reason we check here instead of inside the mssqlcliclient is
                # because we want to raise the Exit exception which will be
                # caught by the try/except block that wraps the mssqlcliclient execute
                # statement.
                if self.quit_command(document.text):
                    raise EOFError

                try:
                    document = self.handle_editor_command(self.cli, document)
                except RuntimeError as e:
                    self.logger.error("sql: %r, error: %r", document.text, e)
                    self.logger.error("traceback: %r", traceback.format_exc())
                    click.secho(str(e), err=True, fg='red')
                    continue

                # Initialize default metaquery in case execution fails
                query = MetaQuery(query=document.text, successful=False)
                query = self.execute_command(document.text, query)
                self.now = dt.datetime.today()

                if not query.contains_secure_statement:
                    # Allow MssqlCompleter to learn user's preferred keywords, etc.
                    with self._completer_lock:
                        self.completer.extend_query_history(document.text)

                    self.query_history.append(query)

        except EOFError:
            self.mssqlcliclient_main.shutdown()
            if not self.less_chatty:
                print('Goodbye!')

    def _build_cli(self, history):

        def set_vi_mode(value):
            self.vi_mode = value

        key_binding_manager = mssqlcli_bindings(
            get_vi_mode_enabled=lambda: self.vi_mode,
            set_vi_mode_enabled=set_vi_mode)

        def prompt_tokens(_):
            prompt = self.get_prompt(self.prompt_format)
            return [(Token.Prompt, prompt)]

        def get_continuation_tokens(cli, width):
            continuation = self.multiline_continuation_char * (width - 1) + ' '
            return [(Token.Continuation, continuation)]

        get_toolbar_tokens = create_toolbar_tokens_func(
            lambda: self.vi_mode, None,
            None,
            None)

        layout = create_prompt_layout(
            lexer=PygmentsLexer(PostgresLexer),
            reserve_space_for_menu=self.min_num_menu_lines,
            get_prompt_tokens=prompt_tokens,
            get_continuation_tokens=get_continuation_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            multiline=True,
            extra_input_processors=[
                # Highlight matching brackets while editing.
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
            ])

        with self._completer_lock:
            buf = MssqlBuffer(
                auto_suggest=AutoSuggestFromHistory(),
                always_multiline=self.multi_line,
                multiline_mode=self.multiline_mode,
                completer=self.completer,
                history=history,
                complete_while_typing=Always(),
                accept_action=AcceptAction.RETURN_DOCUMENT)

            editing_mode = EditingMode.VI if self.vi_mode else EditingMode.EMACS

            application = Application(
                style=style_factory(self.syntax_style, self.cli_style),
                layout=layout,
                buffer=buf,
                key_bindings_registry=key_binding_manager.registry,
                on_exit=AbortAction.RAISE_EXCEPTION,
                on_abort=AbortAction.RETRY,
                ignore_case=True,
                editing_mode=editing_mode)

            cli = CommandLineInterface(application=application,
                                       eventloop=self.eventloop)

            return cli

    def _should_show_limit_prompt(self, status, rows):
        """returns True if limit prompt should be shown, False otherwise."""
        if not rows:
            return False
        return self.row_limit > 0 and len(rows) > 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)

        returns (results, MetaQuery)
        """
        all_success = True
        meta_changed = False  # CREATE, ALTER, DROP, etc
        mutated = False  # INSERT, DELETE, etc
        db_changed = False
        contains_secure_statement = False
        path_changed = False
        output = []
        total = 0

        # Run the query.
        start = time()

        # mssql-cli
        if not self.mssqlcliclient_main.connect_to_database():
            click.secho(u'No connection to server. Exiting.')
            exit(1)

        for rows, columns, status, sql, is_error in \
                self.mssqlcliclient_main.execute_query(text):

            total = time() - start
            if self._should_show_limit_prompt(status, rows):
                click.secho('The result set has more than %s rows.'
                            % self.row_limit, fg='red')
                if not click.confirm('Do you want to continue?'):
                    click.secho("Aborted!", err=True, fg='red')
                    break

            contains_secure_statement = security_words_found_in(sql)

            if is_error:
                output.append(status)
                all_success = False
                continue

            if self.auto_expand:
                max_width = self.cli.output.get_size().columns
            else:
                max_width = None

            settings = OutputSettings(
                table_format=self.table_format,
                dcmlfmt=self.decimal_format,
                floatfmt=self.float_format,
                missingval=self.null_string,
                expanded=self.expanded_output,
                max_width=max_width,
                case_function=(
                    self.completer.case if self.settings['case_column_headers']
                    else
                    lambda x: x
                )
            )

            formatted = self.format_output(None, rows, columns, status, settings)
            output.extend(formatted)

            db_changed, new_db_name = self.has_change_db_cmd(sql)

            if new_db_name:
                self.logger.info('Database context changed.')
                self.mssqlcliclient_main.database = new_db_name

            if all_success:
                meta_changed = meta_changed or self.has_meta_cmd(text)

        return output, MetaQuery(
            text, all_success, total, meta_changed, db_changed, path_changed, mutated, contains_secure_statement)

    def _handle_server_closed_connection(self):
        """Used during CLI execution"""
        reconnect = click.prompt(
            'Connection reset. Reconnect (Y/n)',
            show_default=False, type=bool, default=True)
        if reconnect:
            try:
                self.reset()

                click.secho('Reconnected!\nTry the command again.', fg='green')
            except Exception as e:
                click.secho(str(e), err=True, fg='red')

    def reset(self):
        """
        Reset mssqlcli client with a new sql tools service and connection.
        """
        try:
            self.sqltoolsclient.shutdown()
            self.sqltoolsclient = SqlToolsClient()

            self.mssqlcliclient_main = self.mssqlcliclient_main.clone(self.sqltoolsclient)

            if not self.mssqlcliclient_main.connect_to_database():
                click.secho('Unable reconnect to server {0}; database {1}.'.format(
                    self.mssqlcliclient_main.server_name,
                    self.mssqlcliclient_main.database),
                    err=True, fg='yellow')

                self.logger.info(u'Unable to reset connection to server {0}; database {1}'.format(
                    self.mssqlcliclient_main.server_name,
                    self.mssqlcliclient_main.database))
                exit(1)
        except Exception as e:
            self.logger.error(u'Error in reset : {0}'.format(e.message))
            raise e

    def refresh_completions(self, history=None, persist_priorities='all'):
        # Clone mssqlcliclient to create a new connection with a new owner Uri.
        mssqlclclient_completion_refresher = self.mssqlcliclient_main.clone()

        callback = functools.partial(self._on_completions_refreshed,
                                     persist_priorities=persist_priorities)

        self.completion_refresher.refresh(mssqcliclient=mssqlclclient_completion_refresher,
                                          callbacks=callback,
                                          history=history,
                                          settings=self.settings)

        return [(None, None, None,
                 'Auto-completion refresh started in the background.')]

    def _on_completions_refreshed(self, new_completer, persist_priorities):
        self._swap_completer_objects(new_completer, persist_priorities)

        if self.cli:
            # After refreshing, redraw the CLI to clear the statusbar
            # "Refreshing completions..." indicator
            self.cli.request_redraw()

    def _swap_completer_objects(self, new_completer, persist_priorities):
        """Swap the completer object in cli with the newly created completer.

            persist_priorities is a string specifying how the old completer's
            learned prioritizer should be transferred to the new completer.

              'none'     - The new prioritizer is left in a new/clean state

              'all'      - The new prioritizer is updated to exactly reflect
                           the old one

              'keywords' - The new prioritizer is updated with old keyword
                           priorities, but not any other.
        """
        with self._completer_lock:
            old_completer = self.completer
            self.completer = new_completer

            if persist_priorities == 'all':
                # Just swap over the entire prioritizer
                new_completer.prioritizer = old_completer.prioritizer
            elif persist_priorities == 'keywords':
                # Swap over the entire prioritizer, but clear name priorities,
                # leaving learned keyword priorities alone
                new_completer.prioritizer = old_completer.prioritizer
                new_completer.prioritizer.clear_names()
            elif persist_priorities == 'none':
                # Leave the new prioritizer as is
                pass

            # When mssql-cli is first launched we call refresh_completions before
            # instantiating the cli object. So it is necessary to check if cli
            # exists before trying the replace the completer object in cli.
            if self.cli:
                self.cli.current_buffer.completer = new_completer

    def get_completions(self, text, cursor_positition):
        with self._completer_lock:
            return self.completer.get_completions(
                Document(text=text, cursor_position=cursor_positition), None)

    def get_prompt(self, string):
        string = string.replace('\\t', self.now.strftime('%x %X'))
        string = string.replace('\\u', self.mssqlcliclient_main.user_name or '(none)')
        string = string.replace('\\h', self.mssqlcliclient_main.prompt_host or '(none)')
        string = string.replace('\\d', self.mssqlcliclient_main.database or '(none)')
        string = string.replace('\\p', str(self.mssqlcliclient_main.prompt_port) or '(none)')
        string = string.replace('\\n', "\n")
        return string

    def get_last_query(self):
        """Get the last query executed or None."""
        return self.query_history[-1][0] if self.query_history else None

    def has_meta_cmd(self, query):
        """Determines if the completion needs a refresh by checking if the sql
        statement is an alter, create, drop, commit or rollback."""
        try:
            first_token = query.split()[0]
            if first_token.lower() in ('alter', 'create', 'drop'):
                return True
        except Exception:
            return False

        return False

    def has_change_db_cmd(self, query):
        """Determines if the statement is a database switch such as 'use' or '\\c'
           Returns (True, DBName) or (False, None)
        """
        try:
            first_token = query.split()[0]
            if first_token.lower() in ('use', '\\c', '\\connect'):
                return True, query.split()[1].strip('"')
        except Exception:
            return False, None

        return False, None

    def quit_command(self, sql):
        return (sql.strip().lower() == 'exit' or
                sql.strip().lower() == 'quit' or
                sql.strip() == r'\q' or
                sql.strip() == ':q')

    def format_output(self, title, cur, headers, status, settings):
        output = []
        expanded = (settings.expanded or settings.table_format == 'vertical')
        table_format = ('vertical' if settings.expanded else
                        settings.table_format)
        max_width = settings.max_width
        case_function = settings.case_function
        formatter = TabularOutputFormatter(format_name=table_format)

        def format_array(val):
            if val is None:
                return settings.missingval
            if not isinstance(val, list):
                return val
            return '{' + ','.join(text_type(format_array(e)) for e in val) + '}'

        def format_arrays(data, headers, **_):
            data = list(data)
            for row in data:
                row[:] = [
                    format_array(val) if isinstance(val, list) else val
                    for val in row
                ]

            return data, headers

        output_kwargs = {
            'sep_title': 'RECORD {n}',
            'sep_character': '-',
            'sep_length': (1, 25),
            'missing_value': settings.missingval,
            'integer_format': settings.dcmlfmt,
            'float_format': settings.floatfmt,
            'preprocessors': (format_numbers, format_arrays),
            'disable_numparse': True,
            'preserve_whitespace': True
        }
        if not settings.floatfmt:
            output_kwargs['preprocessors'] = (align_decimals, )

        if title:
            output.append(title)

        if cur:
            headers = [case_function(utf8tounicode(x)) for x in headers]
            if max_width is not None:
                cur = list(cur)
            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)
                if isinstance(formatted, text_type):
                    formatted = iter(formatted.splitlines())

            output = itertools.chain(output, formatted)

        if status:  # Only print the status if it's not None.
            output = itertools.chain(output, [status])

        return output
Beispiel #2
0
class MssqlCli(object):
    # pylint: disable=too-many-instance-attributes, useless-object-inheritance

    max_len_prompt = 30
    default_prompt = '\\d> '

    def set_default_pager(self, config):
        configured_pager = config['main'].get('pager')
        os_environ_pager = os.environ.get('PAGER')
        is_less_installed = is_command_valid(['less', '--version'])
        default_pager = configured_pager or os_environ_pager or \
                        ('less -SRXF' if is_less_installed else False) or None

        if configured_pager:
            self.logger.info('Default pager found in config file: "%s"',
                             configured_pager)
        elif os_environ_pager:
            self.logger.info(
                'Default pager found in PAGER environment variable: "%s"',
                os_environ_pager)
        elif is_less_installed:
            self.logger.info('Default pager set to Less')
        else:
            self.logger.info(
                'No default pager found in environment. Using os default pager'
            )

        # Set default set of less recommended options, if they are not already set.
        # They are ignored if pager is different than less.
        if not os.environ.get('LESS'):
            os.environ['LESS'] = '-SRXF'

        if default_pager is not None:
            os.environ['PAGER'] = default_pager
        return default_pager

    def __init__(self, options):

        # Load config.
        c = self.config = get_config(options.mssqlclirc_file)

        self.initialize_logging()
        self.logger = logging.getLogger(u'mssqlcli.main')

        self.interactive_mode = options.interactive_mode

        self.table_format = c['main']['table_format']
        self.decimal_format = c['data_formats']['decimal']
        self.float_format = c['data_formats']['float']
        self.null_string = c['main'].get('null_string', '<null>')
        self.expanded_output = c['main']['expand'] == 'always'
        self.integrated_auth = options.integrated_auth
        self.less_chatty = bool(options.less_chatty) or c['main'].as_bool(
            'less_chatty') or self.interactive_mode

        keyword_casing = c['main']['keyword_casing']
        self.settings = {
            'casing_file': get_casing_file(c),
            'generate_casing_file': c['main'].as_bool('generate_casing_file'),
            'generate_aliases': c['main'].as_bool('generate_aliases'),
            'asterisk_column_order': c['main']['asterisk_column_order'],
            'qualify_columns': c['main']['qualify_columns'],
            'case_column_headers': c['main'].as_bool('case_column_headers'),
            'search_path_filter': c['main'].as_bool('search_path_filter'),
            'single_connection': False,
            'less_chatty': self.less_chatty,
            'keyword_casing': keyword_casing,
        }

        if self.interactive_mode:
            pager = self.set_default_pager(c)
            self.prompt_session = None

            # set auto_expand to false if less is detected with auto expand
            self.auto_expand = options.auto_vertical_output \
                or (c['main']['expand'] == 'auto' and pager != 'less -SRXF')
            self.multiline = c['main'].as_bool('multi_line')
            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)
            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']
            self.cli_style = c['colors']
            self.output_style = style_factory_output(self.syntax_style,
                                                     self.cli_style)
            self.wider_completion_menu = c['main'].as_bool(
                'wider_completion_menu')
            self.on_error = c['main']['on_error'].upper()

            self.now = dt.datetime.today()

            self.completion_refresher = CompletionRefresher()

            self.query_history = []

            # Initialize completer
            smart_completion = c['main'].get('smart_completion',
                                             'True').lower() == 'true'
            self.completer = MssqlCompleter(smart_completion=smart_completion,
                                            settings=self.settings)
            self._completer_lock = threading.Lock()

        # input and output file are for non-interactive mode
        self.input_file = options.input_file
        self.output_file = options.output_file

        self.query = options.query

        self.sqltoolsclient = SqlToolsClient(
            enable_logging=options.enable_sqltoolsservice_logging)
        self.mssqlcliclient_main = MssqlCliClient(options, self.sqltoolsclient)

        # exit and return error if user enters interactive mode with -i or -o arguments enabled
        if self.interactive_mode and (self.input_file or self.output_file):
            raise ValueError(
                "Invalid arguments: -i and -o can only be used in non-interactive "
                "mode.")

        # exit and return error if both query text and an input file are specified
        if self.query and self.input_file:
            raise ValueError(
                "Invalid arguments: either query [-Q] or input file [-i] may be "
                "specified.")

    def __del__(self):
        # Shut-down sqltoolsservice
        if self.sqltoolsclient:
            self.sqltoolsclient.shutdown()

    # TODO: possibly use at a later date for expanded output file functionality
    # def write_to_file(self, pattern, **_):
    #     if not pattern:
    #         self.output_file = None
    #         message = 'File output disabled'
    #         return [(None, None, None, message, '', True)]
    #     filename = os.path.abspath(os.path.expanduser(pattern))
    #     if not os.path.isfile(filename):
    #         try:
    #             open(filename, 'w').close()
    #         except IOError as e:
    #             self.output_file = None
    #             message = str(e) + '\nFile output disabled'
    #             return [(None, None, None, message, '', False)]
    #     self.output_file = filename
    #     message = 'Writing to file "%s"' % self.output_file
    #     return [(None, None, None, message, '', True)]

    def initialize_logging(self):
        log_file = self.config['main']['log_file']
        if log_file == 'default':
            log_file = config_location() + 'mssqlcli.log'
        ensure_dir_exists(log_file)
        log_level = self.config['main']['log_level']

        # Disable logging if value is NONE by switching to a no-op handler.
        # Set log level to a high value so it doesn't even waste cycles getting
        # called.
        if log_level.upper() == 'NONE':
            handler = logging.NullHandler()
        else:
            # creates a log buffer with max size of 20 MB and 5 backup files
            handler = RotatingFileHandler(os.path.expanduser(log_file),
                                          encoding='utf-8',
                                          maxBytes=1024 * 1024 * 20,
                                          backupCount=5)

        level_map = {
            'CRITICAL': logging.CRITICAL,
            'ERROR': logging.ERROR,
            'WARNING': logging.WARNING,
            'INFO': logging.INFO,
            'DEBUG': logging.DEBUG,
            'NONE': logging.CRITICAL
        }

        log_level = level_map[log_level.upper()]

        formatter = logging.Formatter(
            '%(asctime)s (%(process)d/%(threadName)s) '
            '%(name)s %(levelname)s - %(message)s')

        handler.setFormatter(formatter)

        root_logger = logging.getLogger('mssqlcli')
        root_logger.addHandler(handler)
        root_logger.setLevel(log_level)

        root_logger.info('Initializing mssqlcli logging.')
        root_logger.debug('Log file %r.', log_file)

    def set_main_mssqlcli_client(self, mssqlcli_client):
        self.mssqlcliclient_main = mssqlcli_client

    def connect_to_database(self):
        owner_uri, error_messages = self.mssqlcliclient_main.connect_to_database(
        )
        if not owner_uri and error_messages:
            click.secho('\n'.join(error_messages), err=True, fg='yellow')
            self.logger.debug('Database connection failed: %r.',
                              error_messages)
            sys.exit(1)

    def handle_editor_command(self, text):
        r"""
        Editor command is any query that is prefixed or suffixed
        by a '\e'. The reason for a while loop is because a user
        might edit a query multiple times.
        For eg:
        "select * from \e"<enter> to edit it in vim, then come
        back to the prompt with the edited query "select * from
        blah where q = 'abc'\e" to edit it again.
        :param text: Document
        :return: Document
        """
        # FIXME: using application.pre_run_callables like this here is not the best solution.
        # It's internal api of prompt_toolkit that may change. This was added to fix #668.
        # We may find a better way to do it in the future.
        # pylint: disable=no-member
        editor_command = special.editor_command(text)
        while editor_command:
            filename = special.get_filename(text)
            query = (special.get_editor_query(text) or self.get_last_query())
            sql, message = special.open_external_editor(filename, sql=query)
            if message:
                # Something went wrong. Raise an exception and bail.
                raise RuntimeError(message)
            while True:
                try:
                    text = self.prompt_session.prompt(default=sql)
                    break
                except KeyboardInterrupt:
                    sql = ""

            editor_command = special.editor_command(text)
        return text

    def _execute_interactive_command(self, text):
        """ Runs commands in the interactive CLI mode. """

        logger = self.logger

        # Initialize default metaquery in case execution fails
        query = MetaQuery(query=text, successful=False)

        try:
            output, query = self._evaluate_command(text)
        except KeyboardInterrupt:
            # Issue where Ctrl+C propagates to sql tools service process and kills it,
            # so that query/cancel request can't be sent.
            # Right now the sql_tools_service process is killed and we restart
            # it with a new connection.
            click.secho(u'Cancelling query...', err=True, fg='red')
            self.reset()
            logger.debug("cancelled query, sql: %r", text)
            click.secho("Query cancelled.", err=True, fg='red')

        except NotImplementedError:
            click.secho('Not Yet Implemented.', fg="yellow")
        else:
            if query.total_time > 1:
                # pylint: disable=no-member
                print('Time: %0.03fs (%s)' %
                      (query.total_time,
                       humanize.time.naturaldelta(query.total_time)))
            else:
                print('Time: %0.03fs' % query.total_time)

            # Check if we need to update completions, in order of most
            # to least drastic changes
            if query.db_changed:
                with self._completer_lock:
                    self.completer.reset_completions()
                self.refresh_completions(persist_priorities='keywords')
            elif query.meta_changed:
                self.refresh_completions(persist_priorities='all')

        if not query.contains_secure_statement:
            # Allow MssqlCompleter to learn user's preferred keywords, etc.
            with self._completer_lock:
                self.completer.extend_query_history(text)

            self.query_history.append(query)

        return output

    def execute_query(self, text):
        """ Processes a query string and outputs to file or terminal """
        if self.interactive_mode:
            output = self._execute_interactive_command(text)
        else:
            # non-interactive mode
            output, _ = self._evaluate_command(text)

        self._output_query(output)
        return output

    def _output_query(self, output):
        """ Specifies how query output is handled """
        if self.interactive_mode:
            click.echo_via_pager('\n'.join(output))
        else:
            if self.output_file:
                try:
                    with open(self.output_file, 'w', encoding='utf-8') as f:
                        click.echo('\n'.join(output), file=f)
                except IOError as e:
                    click.secho(str(e), err=True, fg='red')
                    sys.exit(1)
            else:
                click.echo('\n'.join(output))

    def run(self):
        """ Spins up CLI. """

        # raise error if interactive mode is set to false here
        if not self.interactive_mode:
            raise ValueError(
                "Invalid arguments: 'run' must be used in interactive mode! Please set "
                "interactive_mode to True.")

        # exit and return error if user enters interactive mode with -o argument enabled
        if self.output_file:
            raise ValueError(
                "Invalid arguments: -o must be used with interactive mode set to "
                "false.")

        history_file = self.config['main']['history_file']
        if history_file == 'default':
            history_file = config_location() + 'history'
        history = MssqlFileHistory(os.path.expanduser(history_file))

        self.refresh_completions(history=history, persist_priorities='none')

        self.prompt_session = self._build_cli(history)

        if not self.less_chatty:
            print('Version: {}'.format(__version__))
            print('Mail: [email protected]')
            print('Home: http://github.com/dbcli/mssql-cli')

        try:
            while True:
                try:
                    text = self.prompt_session.prompt()
                except KeyboardInterrupt:
                    continue

                # The reason we check here instead of inside the mssqlcliclient is
                # because we want to raise the Exit exception which will be
                # caught by the try/except block that wraps the mssqlcliclient execute
                # statement.
                if self.quit_command(text):
                    raise EOFError

                try:
                    text = self.handle_editor_command(text)
                except RuntimeError as e:
                    self.logger.error("sql: %r, error: %r", text, e)
                    self.logger.error("traceback: %r", traceback.format_exc())
                    click.secho(str(e), err=True, fg='red')
                    continue

                self.execute_query(text)
                self.now = dt.datetime.today()

        except EOFError:
            self.mssqlcliclient_main.shutdown()
            if not self.less_chatty:
                print(localized.goodbye())

    def _build_cli(self, history):
        """
            Builds prompt session.
            NOTE: PROMPT-SESSION USES THIS AS DEPENDENCY.
        """
        def get_message():
            prompt = self.get_prompt(self.prompt_format)
            return [(u'class:prompt', prompt)]

        def get_continuation(width, line_number, is_soft_wrap):
            """
                NOTE: updating parameters will cause prompt session to crash.
            """
            # pylint: disable=unused-argument
            continuation = self.multiline_continuation_char * (width - 1) + ' '
            return [(u'class:continuation', continuation)]

        get_toolbar_tokens = create_toolbar_tokens_func(self)

        if self.wider_completion_menu:
            complete_style = CompleteStyle.MULTI_COLUMN
        else:
            complete_style = CompleteStyle.COLUMN

        with self._completer_lock:
            self.prompt_session = PromptSession(
                message=get_message,
                style=style_factory(self.syntax_style, self.cli_style),

                # Layout options.
                lexer=PygmentsLexer(PostgresLexer),
                prompt_continuation=get_continuation,
                bottom_toolbar=get_toolbar_tokens,
                complete_style=complete_style,
                input_processors=[
                    ConditionalProcessor(
                        processor=HighlightMatchingBracketProcessor(
                            chars='[](){}'),
                        #pylint: disable=invalid-unary-operand-type
                        filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
                    # Render \t as 4 spaces instead of "^I"
                    TabsProcessor(char1=u' ', char2=u' ')
                ],
                reserve_space_for_menu=self.min_num_menu_lines,

                # Buffer options.
                multiline=mssql_is_multiline(self),
                completer=ThreadedCompleter(
                    DynamicCompleter(lambda: self.completer)),
                history=history,
                auto_suggest=AutoSuggestFromHistory(),
                complete_while_typing=True,

                # Key bindings.
                enable_system_prompt=True,
                enable_open_in_editor=True,

                # Other options.
                key_bindings=mssqlcli_bindings(self),
                editing_mode=EditingMode.VI
                if self.vi_mode else EditingMode.EMACS,
                search_ignore_case=True)

            return self.prompt_session

    def _should_show_limit_prompt(self, status, rows):
        """
            Returns True if limit prompt should be shown, False otherwise.
            NOTE: updating parameters will cause prompt session to crash.
        """
        # pylint: disable=unused-argument
        if not rows:
            return False
        return self.interactive_mode and self.row_limit > 0 and len(
            rows) > 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)

            returns (results, MetaQuery)
        """
        # pylint: disable=too-many-locals

        all_success = True
        meta_changed = False  # CREATE, ALTER, DROP, etc
        mutated = False  # INSERT, DELETE, etc
        db_changed = False
        contains_secure_statement = False
        path_changed = False
        output = []
        total = 0

        # Run the query.
        start = time()

        # mssql-cli
        if not self.mssqlcliclient_main.connect_to_database():
            click.secho(u'No connection to server. Exiting.')
            sys.exit(1)

        for rows, columns, status, sql, is_error in \
                self.mssqlcliclient_main.execute_query(text):

            total = time() - start
            if self._should_show_limit_prompt(status, rows):
                click.secho('The result set has more than %s rows.' %
                            self.row_limit,
                            fg='red')
                if not click.confirm('Do you want to continue?'):
                    click.secho("Aborted!", err=True, fg='red')
                    break

            contains_secure_statement = security_words_found_in(sql)

            if is_error:
                output.append(status)
                all_success = False
                continue

            if self.interactive_mode and self.auto_expand and self.prompt_session:
                max_width = self.prompt_session.output.get_size().columns
            else:
                max_width = None

            settings = OutputSettings(
                table_format=self.table_format,
                dcmlfmt=self.decimal_format,
                floatfmt=self.float_format,
                missingval=self.null_string,
                expanded=self.expanded_output,
                max_width=max_width,
                case_function=(self.completer.case if self.interactive_mode
                               and self.settings['case_column_headers'] else
                               lambda x: x))

            formatted = self.format_output(None, rows, columns, status,
                                           settings)
            output.extend(formatted)

            db_changed, new_db_name = self.has_change_db_cmd(sql)

            if new_db_name:
                self.logger.info('Database context changed.')
                self.mssqlcliclient_main.connected_database = new_db_name

            if all_success:
                meta_changed = meta_changed or self.has_meta_cmd(text)

        return output, MetaQuery(text, all_success, total, meta_changed,
                                 db_changed, path_changed, mutated,
                                 contains_secure_statement)

    def _handle_server_closed_connection(self):
        """Used during CLI execution"""
        reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
                                 show_default=False,
                                 type=bool,
                                 default=True)
        if reconnect:
            self.reset()
            click.secho('Reconnected!\nTry the command again.', fg='green')

    def shutdown(self):
        """ API for shutting down client """
        self.mssqlcliclient_main.shutdown()

    def reset(self):
        """
        Reset mssqlcli client with a new sql tools service and connection.
        """
        self.sqltoolsclient.shutdown()
        self.sqltoolsclient = SqlToolsClient()

        self.mssqlcliclient_main = self.mssqlcliclient_main.clone(
            self.sqltoolsclient)

        database_response = self.mssqlcliclient_main.connect_to_database()
        if not database_response:
            click.secho('Unable reconnect to server %s; database %s.' %
                        (self.mssqlcliclient_main.server_name,
                         self.mssqlcliclient_main.connected_database),
                        err=True,
                        fg='yellow')

            self.logger.info(
                u'Unable to reset connection to server %s; database %s',
                self.mssqlcliclient_main.server_name,
                self.mssqlcliclient_main.connected_database)

            sys.exit(1)
        else:
            owner_uri, error_messages = database_response
            if not owner_uri and error_messages:
                # can occur if database credentials change during reset
                self.logger.error(u'Error in reset : %s', error_messages)
                raise ConnectionResetError(error_messages)

    def refresh_completions(self, history=None, persist_priorities='all'):
        # Clone mssqlcliclient to create a new connection with a new owner Uri.
        mssqlclclient_completion_refresher = self.mssqlcliclient_main.clone()

        callback = functools.partial(self._on_completions_refreshed,
                                     persist_priorities=persist_priorities)

        self.completion_refresher.refresh(
            mssqcliclient=mssqlclclient_completion_refresher,
            callbacks=callback,
            history=history,
            settings=self.settings)

        return [(None, None, None,
                 'Auto-completion refresh started in the background.')]

    def _on_completions_refreshed(self, new_completer, persist_priorities):
        self._swap_completer_objects(new_completer, persist_priorities)

        if self.prompt_session:
            # After refreshing, redraw the CLI to clear the statusbar
            # "Refreshing completions..." indicator
            self.prompt_session.app.invalidate()

    def _swap_completer_objects(self, new_completer, persist_priorities):
        """Swap the completer object with the newly created completer.

            persist_priorities is a string specifying how the old completer's
            learned prioritizer should be transferred to the new completer.

              'none'     - The new prioritizer is left in a new/clean state

              'all'      - The new prioritizer is updated to exactly reflect
                           the old one

              'keywords' - The new prioritizer is updated with old keyword
                           priorities, but not any other.
        """
        with self._completer_lock:
            old_completer = self.completer
            self.completer = new_completer

            if persist_priorities == 'all':
                # Just swap over the entire prioritizer
                new_completer.prioritizer = old_completer.prioritizer
            elif persist_priorities == 'keywords':
                # Swap over the entire prioritizer, but clear name priorities,
                # leaving learned keyword priorities alone
                new_completer.prioritizer = old_completer.prioritizer
                new_completer.prioritizer.clear_names()
            elif persist_priorities == 'none':
                # Leave the new prioritizer as is
                pass

            # When mssql-cli is first launched we call refresh_completions before
            # instantiating the cli object. So it is necessary to check if cli
            # exists before trying the replace the completer object in cli.

            self.completer = new_completer

    def get_completions(self, text, cursor_position):
        with self._completer_lock:
            return self.completer.get_completions(
                Document(text=text, cursor_position=cursor_position), None)

    def get_prompt(self, string):
        string = string.replace('\\t', self.now.strftime('%x %X'))
        string = string.replace('\\u', self.mssqlcliclient_main.user_name
                                or '(none)')
        string = string.replace(
            '\\h', self.mssqlcliclient_main.prompt_host or '(none)')
        string = string.replace(
            '\\d', self.mssqlcliclient_main.connected_database or '(none)')
        string = string.replace(
            '\\p',
            str(self.mssqlcliclient_main.prompt_port) or '(none)')
        string = string.replace('\\n', "\n")
        return string

    def get_last_query(self):
        """Get the last query executed or None."""
        return self.query_history[-1][0] if self.query_history else None

    @staticmethod
    def has_meta_cmd(query):
        """Determines if the completion needs a refresh by checking if the sql
        statement is an alter, create, drop, commit or rollback."""
        if query and isinstance(query, str):
            first_token = query.split()[0]
            if first_token.lower() in ('alter', 'create', 'drop'):
                return True
        return False

    @staticmethod
    def has_change_db_cmd(query):
        """Determines if the statement is a database switch such as 'use' or '\\c'
           Returns (True, DBName) or (False, None)
        """
        if query and isinstance(query, str):
            first_token = query.split()[0]
            if first_token.lower() in ('use', '\\c', '\\connect'):
                return True, query.split()[1].strip('"')

        return False, None

    @staticmethod
    def quit_command(sql):
        return (sql.strip().lower() == 'exit' or sql.strip().lower() == 'quit'
                or sql.strip() == r'\q' or sql.strip() == ':q')

    @staticmethod
    def format_output(title, cur, headers, status, settings):
        # pylint: disable=too-many-locals

        output = []
        expanded = (settings.expanded or settings.table_format == 'vertical')
        table_format = ('vertical'
                        if settings.expanded else settings.table_format)
        max_width = settings.max_width
        case_function = settings.case_function
        formatter = TabularOutputFormatter(format_name=table_format)

        def format_array(val):
            if val is None:
                return settings.missingval
            if not isinstance(val, list):
                return val
            return '{' + ','.join(text_type(format_array(e))
                                  for e in val) + '}'

        def format_arrays(data, headers, **_):
            data = list(data)
            for row in data:
                row[:] = [
                    format_array(val) if isinstance(val, list) else val
                    for val in row
                ]

            return data, headers

        output_kwargs = {
            'sep_title': 'RECORD {n}',
            'sep_character': '-',
            'sep_length': (1, 25),
            'missing_value': settings.missingval,
            'integer_format': settings.dcmlfmt,
            'float_format': settings.floatfmt,
            'preprocessors': (format_numbers, format_arrays),
            'disable_numparse': True,
            'preserve_whitespace': True
        }
        if not settings.floatfmt:
            output_kwargs['preprocessors'] = (align_decimals, )

        if title:
            output.append(title)

        if cur:
            headers = [case_function(x) for x in headers]
            if max_width is not None:
                cur = list(cur)
            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)
                if isinstance(formatted, text_type):
                    formatted = iter(formatted.splitlines())

            output = itertools.chain(output, formatted)

        if status:  # Only print the status if it's not None.
            output = itertools.chain(output, [status])

        return output