def run_cli(self): pgexecute = self.pgexecute logger = self.logger original_less_opts = self.adjust_less_opts() completer = self.completer self.refresh_completions() def set_vi_mode(value): self.vi_mode = value key_binding_manager = pgcli_bindings( get_vi_mode_enabled=lambda: self.vi_mode, set_vi_mode_enabled=set_vi_mode) print('Version:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') def prompt_tokens(cli): return [(Token.Prompt, '%s> ' % pgexecute.dbname)] get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.vi_mode) layout = create_default_layout( lexer=PostgresLexer, reserve_space_for_menu=True, get_prompt_tokens=prompt_tokens, get_bottom_toolbar_tokens=get_toolbar_tokens, extra_input_processors=[ # Highlight matching brackets while editing. ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()), ]) history_file = self.config['main']['history_file'] buf = PGBuffer(always_multiline=self.multi_line, completer=completer, history=FileHistory(os.path.expanduser(history_file)), complete_while_typing=Always()) application = Application( style=style_factory(self.syntax_style), layout=layout, buffer=buf, key_bindings_registry=key_binding_manager.registry, on_exit=AbortAction.RAISE_EXCEPTION) cli = CommandLineInterface(application=application, eventloop=create_eventloop()) try: while True: document = cli.run() # The reason we check here instead of inside the pgexecute is # because we want to raise the Exit exception which will be # caught by the try/except block that wraps the pgexecute.run() # statement. if quit_command(document.text): raise EOFError try: document = self.handle_editor_command(cli, document) except RuntimeError as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') continue # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug('sql: %r', document.text) successful = False # Initialized to [] because res might never get initialized # if an exception occurs in pgexecute.run(). Which causes # finally clause to fail. res = [] start = time() # Run the query. res = pgexecute.run(document.text, self.pgspecial) duration = time() - start successful = True output = [] total = 0 for title, cur, headers, status in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) start = time() threshold = 1000 if (is_select(status) and cur and cur.rowcount > threshold): 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 formatted = format_output( title, cur, headers, status, self.table_format, self.pgspecial.expanded_output) output.extend(formatted) end = time() total += end - start mutating = mutating or is_mutating(status) except KeyboardInterrupt: # Restart connection to the database pgexecute.connect() logger.debug("cancelled query, sql: %r", document.text) click.secho("cancelled query", err=True, fg='red') except NotImplementedError: click.secho('Not Yet Implemented.', fg="yellow") except OperationalError as e: reconnect = True if ('server closed the connection' in utf8tounicode(e.args[0])): reconnect = click.prompt( 'Connection reset. Reconnect (Y/n)', show_default=False, type=bool, default=True) if reconnect: try: pgexecute.connect() click.secho( 'Reconnected!\nTry the command again.', fg='green') except OperationalError as e: click.secho(str(e), err=True, fg='red') else: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') except Exception as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') else: click.echo_via_pager('\n'.join(output)) if self.pgspecial.timing_enabled: print('Command Time: %0.03fs' % duration) print('Format Time: %0.03fs' % total) # Refresh the table names and column names if necessary. if need_completion_refresh(document.text): self.refresh_completions() # Refresh search_path to set default schema. if need_search_path_refresh(document.text): logger.debug('Refreshing search path') completer.set_search_path(pgexecute.search_path()) logger.debug('Search path: %r', completer.search_path) query = Query(document.text, successful, mutating) self.query_history.append(query) except EOFError: print('Goodbye!') finally: # Reset the less opts back to original. logger.debug('Restoring env var LESS to %r.', original_less_opts) os.environ['LESS'] = original_less_opts
def _build_cli(self, history): def set_vi_mode(value): self.vi_mode = value key_binding_manager = pgcli_bindings( get_vi_mode_enabled=lambda: self.vi_mode, set_vi_mode_enabled=set_vi_mode) def prompt_tokens(_): if self.dsn_alias and self.prompt_dsn_format is not None: prompt_format = self.prompt_dsn_format else: prompt_format = self.prompt_format prompt = self.get_prompt(prompt_format) if (prompt_format == self.default_prompt and len(prompt) > self.max_len_prompt): prompt = self.get_prompt('\\d> ') 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, self.completion_refresher.is_refreshing, self.pgexecute.failed_transaction, self.pgexecute.valid_transaction) 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 = PGBuffer( 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 run_cli(self): iterations = 0 sqlexecute = self.sqlexecute logger = self.logger self.configure_pager() if self.smart_completion: self.refresh_completions() author_file = os.path.join(PACKAGE_ROOT, 'AUTHORS') sponsor_file = os.path.join(PACKAGE_ROOT, 'SPONSORS') key_binding_manager = mycli_bindings() if not self.less_chatty: print('Version:', __version__) print('Chat: https://gitter.im/dbcli/mycli') print('Mail: https://groups.google.com/forum/#!forum/mycli-users') print('Home: http://mycli.net') print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file])) def prompt_tokens(cli): prompt = self.get_prompt(self.prompt_format) if self.prompt_format == self.default_prompt and len( prompt) > self.max_len_prompt: prompt = self.get_prompt('\\d> ') return [(Token.Prompt, prompt)] def get_continuation_tokens(cli, width): continuation_prompt = self.get_prompt( self.prompt_continuation_format) return [(Token.Continuation, ' ' * (width - len(continuation_prompt)) + continuation_prompt)] def show_suggestion_tip(): return iterations < 2 def one_iteration(document=None): if document is None: document = self.cli.run() special.set_expanded_output(False) try: document = self.handle_editor_command(self.cli, document) except RuntimeError as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') return if not document.text.strip(): return if self.destructive_warning: destroy = confirm_destructive_query(document.text) if destroy is None: pass # Query was not destructive. Nothing to do here. elif destroy is True: self.echo('Your call!') else: self.echo('Wise choice!') return # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug('sql: %r', document.text) special.write_tee( self.get_prompt(self.prompt_format) + document.text) if self.logfile: self.logfile.write('\n# %s\n' % datetime.now()) self.logfile.write(document.text) self.logfile.write('\n') successful = False start = time() res = sqlexecute.run(document.text) successful = True result_count = 0 for title, cur, headers, status in res: 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): self.echo( 'The result set has more than {} rows.'.format( threshold), fg='red') if not click.confirm('Do you want to continue?'): self.echo("Aborted!", err=True, fg='red') break if self.auto_vertical_output: max_width = self.cli.output.get_size().columns else: max_width = None formatted = self.format_output( title, cur, headers, special.is_expanded_output(), max_width) t = time() - start try: if result_count > 0: self.echo('') try: self.output(formatted, status) except KeyboardInterrupt: pass if special.is_timing_enabled(): self.echo('Time: %0.03fs' % t) except KeyboardInterrupt: pass start = time() result_count += 1 mutating = mutating or is_mutating(status) special.unset_once_if_written() except EOFError as e: raise e except KeyboardInterrupt: # get last connection id connection_id_to_kill = sqlexecute.connection_id logger.debug("connection id to kill: %r", connection_id_to_kill) # Restart connection to the database sqlexecute.connect() try: for title, cur, headers, status in sqlexecute.run( 'kill %s' % connection_id_to_kill): status_str = str(status).lower() if status_str.find('ok') > -1: logger.debug( "cancelled query, connection id: %r, sql: %r", connection_id_to_kill, document.text) self.echo("cancelled query", err=True, fg='red') except Exception as e: self.echo( 'Encountered error while cancelling query: {}'.format( e), err=True, fg='red') except NotImplementedError: self.echo('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.debug("Exception: %r", e) if (e.args[0] in (2003, 2006, 2013)): logger.debug('Attempting to reconnect.') self.echo('Reconnecting...', fg='yellow') try: sqlexecute.connect() logger.debug('Reconnected successfully.') one_iteration(document) return # OK to just return, cuz the recursion call runs to the end. except OperationalError as e: logger.debug('Reconnect failed. e: %r', e) self.echo(str(e), err=True, fg='red') # If reconnection failed, don't proceed further. return else: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') except Exception as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') else: if is_dropping_database(document.text, self.sqlexecute.dbname): self.sqlexecute.dbname = None self.sqlexecute.connect() # Refresh the table names and column names if necessary. if need_completion_refresh(document.text): self.refresh_completions( reset=need_completion_reset(document.text)) finally: if self.logfile is False: self.echo("Warning: This query was not logged.", err=True, fg='red') query = Query(document.text, successful, mutating) self.query_history.append(query) get_toolbar_tokens = create_toolbar_tokens_func( self.completion_refresher.is_refreshing, show_suggestion_tip) layout = create_prompt_layout( lexer=MyCliLexer, multiline=True, 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, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()) ], reserve_space_for_menu=self.get_reserved_space()) with self._completer_lock: buf = CLIBuffer(always_multiline=self.multi_line, completer=self.completer, history=FileHistory( os.path.expanduser( os.environ.get('MYCLI_HISTFILE', '~/.mycli-history'))), auto_suggest=AutoSuggestFromHistory(), complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT) if self.key_bindings == 'vi': editing_mode = EditingMode.VI else: editing_mode = EditingMode.EMACS application = Application( style=style_from_pygments(style_cls=self.output_style), layout=layout, buffer=buf, key_bindings_registry=key_binding_manager.registry, on_exit=AbortAction.RAISE_EXCEPTION, on_abort=AbortAction.RETRY, editing_mode=editing_mode, ignore_case=True) self.cli = CommandLineInterface(application=application, eventloop=create_eventloop()) try: while True: one_iteration() iterations += 1 except EOFError: special.close_tee() if not self.less_chatty: self.echo('Goodbye!')
def run_cli(self): sqlexecute = self.sqlexecute logger = self.logger original_less_opts = self.adjust_less_opts() self.set_pager_from_config() self.initialize_completions() completer = self.completer def set_key_bindings(value): if value not in ('emacs', 'vi'): value = 'emacs' self.key_bindings = value project_root = os.path.dirname(PACKAGE_ROOT) author_file = os.path.join(project_root, 'AUTHORS') sponsor_file = os.path.join(project_root, 'SPONSORS') key_binding_manager = mycli_bindings( get_key_bindings=lambda: self.key_bindings, set_key_bindings=set_key_bindings) print('Version:', __version__) print('Chat: https://gitter.im/dbcli/mycli') print('Mail: https://groups.google.com/forum/#!forum/mycli-users') print('Home: http://mycli.net') print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file])) def prompt_tokens(cli): return [(Token.Prompt, self.get_prompt(self.prompt_format))] get_toolbar_tokens = create_toolbar_tokens_func( lambda: self.key_bindings) layout = create_default_layout( lexer=MyCliLexer, reserve_space_for_menu=True, multiline=True, get_prompt_tokens=prompt_tokens, get_bottom_toolbar_tokens=get_toolbar_tokens, display_completions_in_columns=self.wider_completion_menu, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()), ]) buf = CLIBuffer(always_multiline=self.multi_line, completer=completer, history=FileHistory( os.path.expanduser('~/.mycli-history')), complete_while_typing=Always()) 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, ignore_case=True) cli = CommandLineInterface(application=application, eventloop=create_eventloop()) try: while True: document = cli.run() special.set_expanded_output(False) # The reason we check here instead of inside the sqlexecute is # because we want to raise the Exit exception which will be # caught by the try/except block that wraps the # sqlexecute.run() statement. if quit_command(document.text): raise EOFError try: document = self.handle_editor_command(cli, document) except RuntimeError as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') continue if self.destructive_warning: destroy = confirm_destructive_query(document.text) if destroy is None: pass # Query was not destructive. Nothing to do here. elif destroy is True: self.output('Your call!') else: self.output('Wise choice!') continue # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug('sql: %r', document.text) if self.logfile: self.logfile.write('\n# %s\n' % datetime.now()) self.logfile.write(document.text) self.logfile.write('\n') successful = False start = time() res = sqlexecute.run(document.text) duration = time() - start successful = True output = [] total = 0 for title, cur, headers, status in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) start = time() threshold = 1000 if (is_select(status) and cur and cur.rowcount > threshold): self.output( 'The result set has more than %s rows.' % threshold, fg='red') if not click.confirm('Do you want to continue?'): self.output("Aborted!", err=True, fg='red') break output.extend( format_output(title, cur, headers, status, self.table_format)) end = time() total += end - start mutating = mutating or is_mutating(status) except KeyboardInterrupt: # Restart connection to the database sqlexecute.connect() logger.debug("cancelled query, sql: %r", document.text) self.output("cancelled query", err=True, fg='red') except NotImplementedError: self.output('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.debug("Exception: %r", e) reconnect = True if (e.args[0] in (2003, 2006, 2013)): reconnect = click.prompt( 'Connection reset. Reconnect (Y/n)', show_default=False, type=bool, default=True) if reconnect: logger.debug('Attempting to reconnect.') try: sqlexecute.connect() logger.debug('Reconnected successfully.') self.output( 'Reconnected!\nTry the command again.', fg='green') except OperationalError as e: logger.debug('Reconnect failed. e: %r', e) self.output(str(e), err=True, fg='red') continue # If reconnection failed, don't proceed further. else: # If user chooses not to reconnect, don't proceed further. continue else: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') except Exception as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') else: try: self.output_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if special.is_timing_enabled(): self.output('Command Time: %0.03fs' % duration) self.output('Format Time: %0.03fs' % total) # Refresh the table names and column names if necessary. if need_completion_refresh(document.text): self.refresh_dynamic_completions() query = Query(document.text, successful, mutating) self.query_history.append(query) except EOFError: self.output('Goodbye!') finally: # Reset the less opts back to original. logger.debug('Restoring env var LESS to %r.', original_less_opts) os.environ['LESS'] = original_less_opts os.environ['PAGER'] = special.get_original_pager()
def init_prompt_toolkit_cli(self): if self.simple_prompt or ('JUPYTER_CONSOLE_TEST' in os.environ): # Simple restricted interface for tests so we can find prompts with # pexpect. Multi-line input not supported. @asyncio.coroutine def prompt(): prompt = 'In [%d]: ' % self.execution_count raw = yield from async_input(prompt) return raw self.prompt_for_code = prompt self.print_out_prompt = \ lambda: print('Out[%d]: ' % self.execution_count, end='') return kb = KeyBindings() insert_mode = vi_insert_mode | emacs_insert_mode @kb.add("enter", filter=(has_focus(DEFAULT_BUFFER) & ~has_selection & insert_mode)) def _(event): b = event.current_buffer d = b.document if not (d.on_last_line or d.cursor_position_row >= d.line_count - d.empty_line_count_at_the_end()): b.newline() return # Pressing enter flushes any pending display. This also ensures # the displayed execution_count is correct. self.handle_iopub() more, indent = self.check_complete(d.text) if (not more) and b.accept_handler: b.validate_and_handle() else: b.insert_text('\n' + indent) @kb.add("c-c", filter=has_focus(DEFAULT_BUFFER)) def _(event): event.current_buffer.reset() @kb.add("c-\\", filter=has_focus(DEFAULT_BUFFER)) def _(event): raise EOFError @kb.add("c-z", filter=Condition(lambda: suspend_to_background_supported())) def _(event): event.cli.suspend_to_background() # Pre-populate history from IPython's history database history = InMemoryHistory() last_cell = u"" for _, _, cell in self.history_manager.get_tail( self.history_load_length, include_latest=True): # Ignore blank lines and consecutive duplicates cell = cell.rstrip() if cell and (cell != last_cell): history.append_string(cell) style_overrides = { Token.Prompt: '#009900', Token.PromptNum: '#00ff00 bold', Token.OutPrompt: '#ff2200', Token.OutPromptNum: '#ff0000 bold', Token.RemotePrompt: '#999900', } if self.highlighting_style: style_cls = get_style_by_name(self.highlighting_style) else: style_cls = get_style_by_name('default') # The default theme needs to be visible on both a dark background # and a light background, because we can't tell what the terminal # looks like. These tweaks to the default theme help with that. style_overrides.update({ Token.Number: '#007700', Token.Operator: 'noinherit', Token.String: '#BB6622', Token.Name.Function: '#2080D0', Token.Name.Class: 'bold #2080D0', Token.Name.Namespace: 'bold #2080D0', }) style_overrides.update(self.highlighting_style_overrides) style = merge_styles([ style_from_pygments_cls(style_cls), style_from_pygments_dict(style_overrides), ]) editing_mode = getattr(EditingMode, self.editing_mode.upper()) langinfo = self.kernel_info.get('language_info', {}) lexer = langinfo.get('pygments_lexer', langinfo.get('name', 'text')) # If enabled in the settings, highlight matching brackets # when the DEFAULT_BUFFER has the focus input_processors = [ ConditionalProcessor( processor=HighlightMatchingBracketProcessor(chars='[](){}'), filter=has_focus(DEFAULT_BUFFER) & ~is_done & Condition(lambda: self.highlight_matching_brackets)) ] # Tell prompt_toolkit to use the asyncio event loop. # Obsolete in prompt_toolkit.v3 if not PTK3: use_asyncio_event_loop() self.pt_cli = PromptSession( message=(lambda: PygmentsTokens(self.get_prompt_tokens())), multiline=True, editing_mode=editing_mode, lexer=PygmentsLexer(get_pygments_lexer(lexer)), prompt_continuation=( lambda width, lineno, is_soft_wrap: PygmentsTokens( self.get_continuation_tokens(width))), key_bindings=kb, history=history, completer=JupyterPTCompleter(self.Completer), enable_history_search=True, style=style, input_processors=input_processors, color_depth=(ColorDepth.TRUE_COLOR if self.true_color else None), )
def run_cli(self): iterations = 0 sqlexecute = self.sqlexecute logger = self.logger self.configure_pager() self.refresh_completions() history_file = config_location() + "history" if dir_path_exists(history_file): history = FileHistory(history_file) else: history = None self.echo( 'Error: Unable to open the history file "{}". ' "Your query history will not be saved.".format(history_file), err=True, fg="red", ) key_bindings = cli_bindings(self) if not self.less_chatty: print("Version:", __version__) print( "Mail: https://groups.google.com/forum/#!forum/litecli-users") print("Github: https://github.com/dbcli/litecli") # print("Home: https://litecli.com") def get_message(): prompt = self.get_prompt(self.prompt_format) if (self.prompt_format == self.default_prompt and len(prompt) > self.max_len_prompt): prompt = self.get_prompt("\\d> ") return [("class:prompt", prompt)] def get_continuation(width, line_number, is_soft_wrap): continuation = " " * (width - 1) + " " return [("class:continuation", continuation)] def show_suggestion_tip(): return iterations < 2 def one_iteration(text=None): if text is None: try: text = self.prompt_app.prompt() except KeyboardInterrupt: return special.set_expanded_output(False) try: text = self.handle_editor_command(text) except RuntimeError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg="red") return if not text.strip(): return if self.destructive_warning: destroy = confirm_destructive_query(text) if destroy is None: pass # Query was not destructive. Nothing to do here. elif destroy is True: self.echo("Your call!") else: self.echo("Wise choice!") return # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug("sql: %r", text) special.write_tee(self.get_prompt(self.prompt_format) + text) if self.logfile: self.logfile.write("\n# %s\n" % datetime.now()) self.logfile.write(text) self.logfile.write("\n") successful = False start = time() res = sqlexecute.run(text) self.formatter.query = text successful = True result_count = 0 for title, cur, headers, status in res: 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: self.echo( "The result set has more than {} rows.".format( threshold), fg="red", ) if not confirm("Do you want to continue?"): self.echo("Aborted!", err=True, fg="red") break if self.auto_vertical_output: max_width = self.prompt_app.output.get_size().columns else: max_width = None formatted = self.format_output( title, cur, headers, special.is_expanded_output(), max_width) t = time() - start try: if result_count > 0: self.echo("") try: self.output(formatted, status) except KeyboardInterrupt: pass self.echo("Time: %0.03fs" % t) except KeyboardInterrupt: pass start = time() result_count += 1 mutating = mutating or is_mutating(status) special.unset_once_if_written() except EOFError as e: raise e except KeyboardInterrupt: # get last connection id connection_id_to_kill = sqlexecute.connection_id logger.debug("connection id to kill: %r", connection_id_to_kill) # Restart connection to the database sqlexecute.connect() try: for title, cur, headers, status in sqlexecute.run( "kill %s" % connection_id_to_kill): status_str = str(status).lower() if status_str.find("ok") > -1: logger.debug( "cancelled query, connection id: %r, sql: %r", connection_id_to_kill, text, ) self.echo("cancelled query", err=True, fg="red") except Exception as e: self.echo( "Encountered error while cancelling query: {}".format( e), err=True, fg="red", ) except NotImplementedError: self.echo("Not Yet Implemented.", fg="yellow") except OperationalError as e: logger.debug("Exception: %r", e) if e.args[0] in (2003, 2006, 2013): logger.debug("Attempting to reconnect.") self.echo("Reconnecting...", fg="yellow") try: sqlexecute.connect() logger.debug("Reconnected successfully.") one_iteration(text) return # OK to just return, cuz the recursion call runs to the end. except OperationalError as e: logger.debug("Reconnect failed. e: %r", e) self.echo(str(e), err=True, fg="red") # If reconnection failed, don't proceed further. return else: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg="red") except Exception as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg="red") else: if is_dropping_database(text, self.sqlexecute.dbname): self.sqlexecute.dbname = None self.sqlexecute.connect() # Refresh the table names and column names if necessary. if need_completion_refresh(text): self.refresh_completions(reset=need_completion_reset(text)) finally: if self.logfile is False: self.echo("Warning: This query was not logged.", err=True, fg="red") query = Query(text, successful, mutating) self.query_history.append(query) get_toolbar_tokens = create_toolbar_tokens_func( self, show_suggestion_tip) if self.wider_completion_menu: complete_style = CompleteStyle.MULTI_COLUMN else: complete_style = CompleteStyle.COLUMN with self._completer_lock: if self.key_bindings == "vi": editing_mode = EditingMode.VI else: editing_mode = EditingMode.EMACS self.prompt_app = PromptSession( lexer=PygmentsLexer(LiteCliLexer), reserve_space_for_menu=self.get_reserved_space(), message=get_message, prompt_continuation=get_continuation, bottom_toolbar=get_toolbar_tokens, complete_style=complete_style, input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars="[](){}"), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone(), ) ], tempfile_suffix=".sql", completer=DynamicCompleter(lambda: self.completer), history=history, auto_suggest=AutoSuggestFromHistory(), complete_while_typing=True, multiline=cli_is_multiline(self), style=style_factory(self.syntax_style, self.cli_style), include_default_pygments_style=False, key_bindings=key_bindings, enable_open_in_editor=True, enable_system_prompt=True, enable_suspend=True, editing_mode=editing_mode, search_ignore_case=True, ) try: while True: one_iteration() iterations += 1 except EOFError: special.close_tee() if not self.less_chatty: self.echo("Goodbye!")
def loop(cmd, history_file): def session_toolbar(cli): return _get_toolbar_tokens(cmd.is_conn_available, cmd.username, cmd.connection.client.active_servers) key_binding_manager = KeyBindingManager( enable_search=True, enable_abort_and_exit_bindings=True, enable_system_bindings=True, enable_open_in_editor=True) bind_keys(key_binding_manager.registry) layout = create_layout( message='cr> ', multiline=True, lexer=SqlLexer, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor(chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()) ], get_bottom_toolbar_tokens=session_toolbar) application = Application( layout=layout, buffer=create_buffer(cmd, history_file), style=PygmentsStyle.from_defaults(pygments_style_cls=CrateStyle), key_bindings_registry=key_binding_manager.registry, editing_mode=_get_editing_mode(), on_exit=AbortAction.RAISE_EXCEPTION, on_abort=AbortAction.RETRY, ) eventloop = create_eventloop() output = create_output() cli = CommandLineInterface(application=application, eventloop=eventloop, output=output) def get_num_columns_override(): return output.get_size().columns cmd.get_num_columns = get_num_columns_override while True: try: doc = cli.run(reset_current_buffer=True) if doc: cmd.process(doc.text) except ProgrammingError as e: if '401' in e.message: username = cmd.username password = cmd.password cmd.username = input('Username: '******'Bye!') return
def _build_cli(self, history): key_binding_manager = cli_bindings() def prompt_tokens(cli): prompt = self.get_prompt(self.prompt) if len(prompt) > self.MAX_LEN_PROMPT: prompt = self.get_prompt('\\r:\\d> ') return [(Token.Prompt, prompt)] def get_continuation_tokens(cli, width): prompt = self.get_prompt(self.prompt_continuation_format) token = (Token.Continuation, ' ' * (width - len(prompt)) + prompt) return [token] def show_suggestion_tip(): return self.iterations < 2 get_toolbar_tokens = create_toolbar_tokens_func( self.completion_refresher.is_refreshing, show_suggestion_tip) layout = create_prompt_layout( lexer=Lexer, multiline=True, get_prompt_tokens=prompt_tokens, get_continuation_tokens=get_continuation_tokens, get_bottom_toolbar_tokens=get_toolbar_tokens, display_completions_in_columns=False, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()) ], reserve_space_for_menu=self.get_reserved_space()) with self._completer_lock: buf = CLIBuffer(always_multiline=self.multi_line, completer=self.completer, history=history, auto_suggest=AutoSuggestFromHistory(), complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT) if self.key_bindings == 'vi': editing_mode = EditingMode.VI else: editing_mode = EditingMode.EMACS application = Application( style=style_from_pygments(style_cls=self.output_style), layout=layout, buffer=buf, key_bindings_registry=key_binding_manager.registry, on_exit=AbortAction.RAISE_EXCEPTION, on_abort=AbortAction.RETRY, editing_mode=editing_mode, ignore_case=True) cli = CommandLineInterface(application=application, eventloop=create_eventloop()) return cli
def _build_cli(self, history): key_bindings = pgcli_bindings(self) def get_message(): if self.dsn_alias and self.prompt_dsn_format is not None: prompt_format = self.prompt_dsn_format else: prompt_format = self.prompt_format prompt = self.get_prompt(prompt_format) if ( prompt_format == self.default_prompt and len(prompt) > self.max_len_prompt ): prompt = self.get_prompt("\\d> ") return [("class:prompt", prompt)] def get_continuation(width, line_number, is_soft_wrap): continuation = self.multiline_continuation_char * (width - 1) + " " return [("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: prompt_app = PromptSession( lexer=PygmentsLexer(PostgresLexer), reserve_space_for_menu=self.min_num_menu_lines, message=get_message, prompt_continuation=get_continuation, bottom_toolbar=get_toolbar_tokens, complete_style=complete_style, input_processors=[ # Highlight matching brackets while editing. ConditionalProcessor( processor=HighlightMatchingBracketProcessor(chars="[](){}"), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone(), ), # Render \t as 4 spaces instead of "^I" TabsProcessor(char1=" ", char2=" "), ], auto_suggest=AutoSuggestFromHistory(), tempfile_suffix=".sql", # N.b. pgcli's multi-line mode controls submit-on-Enter (which # overrides the default behaviour of prompt_toolkit) and is # distinct from prompt_toolkit's multiline mode here, which # controls layout/display of the prompt/buffer multiline=True, history=history, completer=ThreadedCompleter(DynamicCompleter(lambda: self.completer)), complete_while_typing=True, style=style_factory(self.syntax_style, self.cli_style), include_default_pygments_style=False, key_bindings=key_bindings, enable_open_in_editor=True, enable_system_prompt=True, enable_suspend=True, editing_mode=EditingMode.VI if self.vi_mode else EditingMode.EMACS, search_ignore_case=True, ) return prompt_app
def run_cli(self): sqlexecute = self.sqlexecute logger = self.logger self.configure_pager() self.refresh_completions() project_root = os.path.dirname(PACKAGE_ROOT) author_file = os.path.join(project_root, 'AUTHORS') sponsor_file = os.path.join(project_root, 'SPONSORS') key_binding_manager = mycli_bindings() if not self.less_chatty: print('Version:', __version__) print('Chat: https://gitter.im/dbcli/mycli') print('Mail: https://groups.google.com/forum/#!forum/mycli-users') print('Home: http://mycli.net') print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file])) def prompt_tokens(cli): return [(Token.Prompt, self.get_prompt(self.prompt_format))] def get_continuation_tokens(cli, width): continuation_prompt = self.get_prompt( self.prompt_continuation_format) return [(Token.Continuation, ' ' * (width - len(continuation_prompt)) + continuation_prompt)] def one_iteration(document=None): if document is None: document = self.cli.run(reset_current_buffer=True) special.set_expanded_output(False) # The reason we check here instead of inside the sqlexecute is # because we want to raise the Exit exception which will be # caught by the try/except block that wraps the # sqlexecute.run() statement. if quit_command(document.text): raise EOFError try: document = self.handle_editor_command(self.cli, document) except RuntimeError as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') return if self.destructive_warning: destroy = confirm_destructive_query(document.text) if destroy is None: pass # Query was not destructive. Nothing to do here. elif destroy is True: self.output('Your call!') else: self.output('Wise choice!') return # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug('sql: %r', document.text) if self.logfile: self.logfile.write('\n# %s\n' % datetime.now()) self.logfile.write(document.text) self.logfile.write('\n') successful = False start = time() res = sqlexecute.run(document.text) successful = True output = [] total = 0 for title, cur, headers, status in res: 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): self.output('The result set has more than %s rows.' % threshold, fg='red') if not click.confirm('Do you want to continue?'): self.output("Aborted!", err=True, fg='red') break if self.auto_vertical_output: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output(title, cur, headers, status, self.table_format, special.is_expanded_output(), max_width) output.extend(formatted) end = time() total += end - start mutating = mutating or is_mutating(status) except UnicodeDecodeError as e: import pymysql if pymysql.VERSION < (0, 6, 7): message = ( 'You are running an older version of pymysql.\n' 'Please upgrade to 0.6.7 or above to view binary data.\n' 'Try \'pip install -U pymysql\'.') self.output(message) else: raise e except KeyboardInterrupt: # Restart connection to the database sqlexecute.connect() logger.debug("cancelled query, sql: %r", document.text) self.output("cancelled query", err=True, fg='red') except NotImplementedError: self.output('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.debug("Exception: %r", e) if (e.args[0] in (2003, 2006, 2013)): logger.debug('Attempting to reconnect.') self.output('Reconnecting...', fg='yellow') try: sqlexecute.connect() logger.debug('Reconnected successfully.') one_iteration(document) return # OK to just return, cuz the recursion call runs to the end. except OperationalError as e: logger.debug('Reconnect failed. e: %r', e) self.output(str(e), err=True, fg='red') return # If reconnection failed, don't proceed further. else: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') except Exception as e: logger.error("sql: %r, error: %r", document.text, e) logger.error("traceback: %r", traceback.format_exc()) self.output(str(e), err=True, fg='red') else: try: if special.is_pager_enabled(): self.output_via_pager('\n'.join(output)) else: self.output('\n'.join(output)) except KeyboardInterrupt: pass if special.is_timing_enabled(): self.output('Time: %0.03fs' % total) # Refresh the table names and column names if necessary. if need_completion_refresh(document.text): self.refresh_completions( reset=need_completion_reset(document.text)) finally: if self.logfile is False: self.output("Warning: This query was not logged.", err=True, fg='red') query = Query(document.text, successful, mutating) self.query_history.append(query) get_toolbar_tokens = create_toolbar_tokens_func( self.completion_refresher.is_refreshing) layout = create_prompt_layout( lexer=MyCliLexer, multiline=True, 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, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()), ]) with self._completer_lock: buf = CLIBuffer(always_multiline=self.multi_line, completer=self.completer, history=FileHistory( os.path.expanduser( os.environ.get('MYCLI_HISTFILE', '~/.mycli-history'))), complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT) if self.key_bindings == 'vi': editing_mode = EditingMode.VI else: editing_mode = 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, editing_mode=editing_mode, ignore_case=True) self.cli = CommandLineInterface(application=application, eventloop=create_eventloop()) try: while True: one_iteration() except EOFError: if not self.less_chatty: self.output('Goodbye!')
def run_cli(self): """ Run the main loop """ print('Version:', __version__) print('Home: http://wharfee.com') history = FileHistory(os.path.expanduser('~/.wharfee-history')) toolbar_handler = create_toolbar_handler(self.get_long_options, self.get_fuzzy_match) layout = create_prompt_layout( message='wharfee> ', lexer=CommandLexer, get_bottom_toolbar_tokens=toolbar_handler, extra_input_processors=[ ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()) ]) cli_buffer = Buffer(history=history, completer=self.completer, complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT) manager = get_key_manager(self.set_long_options, self.get_long_options, self.set_fuzzy_match, self.get_fuzzy_match) application = Application(style=style_factory(self.theme), layout=layout, buffer=cli_buffer, key_bindings_registry=manager.registry, on_exit=AbortAction.RAISE_EXCEPTION, on_abort=AbortAction.RETRY, ignore_case=True) eventloop = create_eventloop() self.dcli = CommandLineInterface(application=application, eventloop=eventloop) while True: try: document = self.dcli.run(True) self.handler.handle_input(document.text) if isinstance(self.handler.output, GeneratorType): output_stream(self.handler.command, self.handler.output, self.handler.log) elif self.handler.output is not None: lines = format_data(self.handler.command, self.handler.output) click.echo_via_pager('\n'.join(lines)) if self.handler.after: for line in self.handler.after(): click.echo(line) if self.handler.exception: # This was handled, just log it. self.logger.warning('An error was handled: %r', self.handler.exception) self.refresh_completions() except OptionError as ex: self.logger.debug('Error: %r.', ex) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(ex.msg, fg='red') except KeyboardInterrupt: # user pressed Ctrl + C if self.handler.after: click.echo('') for line in self.handler.after(): click.echo(line) self.refresh_completions() except DockerPermissionException as ex: self.logger.debug('Permission exception: %r.', ex) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(ex), fg='red') except EOFError: # exit out of the CLI break # TODO: uncomment for release except Exception as ex: self.logger.debug('Exception: %r.', ex) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(ex), fg='red') self.revert_less_opts() self.write_config_file() print('Goodbye!')
def create_radian_prompt_session(options, settings): local_history_file = settings.local_history_file global_history_file = settings.global_history_file if options.no_history: history = ModalInMemoryHistory() elif not options.global_history and os.path.exists(local_history_file): history = ModalFileHistory(os.path.abspath(local_history_file)) else: history_file = os.path.join(os.path.expanduser(global_history_file)) history_file = os.path.expandvars(history_file) history_file_dir = os.path.dirname(history_file) if not os.path.exists(history_file_dir): os.makedirs(history_file_dir, 0o700) history = ModalFileHistory(history_file) if is_windows(): output = None else: output = CustomOutput.from_pty(sys.stdout, term=get_term_environment_variable()) def get_inputhook(): # make testing more robust if "RADIAN_NO_INPUTHOOK" in os.environ: return None terminal_width = [None] def _(context): output_width = session.app.output.get_size().columns if output_width and terminal_width[0] != output_width: terminal_width[0] = output_width setoption("width", max(terminal_width[0], 20)) while True: if context.input_is_ready(): break try: if peek_event(): with session.app.input.detach(): with session.app.input.rare_mode(): process_events() else: polled_events() except Exception: pass time.sleep(1.0 / 30) return _ def vi_mode_prompt(): if session.editing_mode.lower( ) == "vi" and settings.show_vi_mode_prompt: im = session.app.vi_state.input_mode vi_mode_prompt = settings.vi_mode_prompt if isinstance(vi_mode_prompt, string_types): return vi_mode_prompt.format(str(im)[3:6]) else: return vi_mode_prompt[str(im)[3:6]] return "" def message(): if hasattr(session.current_mode, "get_message"): return ANSI(vi_mode_prompt() + session.current_mode.get_message()) elif hasattr(session.current_mode, "message"): message = session.current_mode.message if callable(message): return ANSI(vi_mode_prompt() + message()) else: return ANSI(vi_mode_prompt() + message) else: return "" session = ModalPromptSession( message=message, color_depth=ColorDepth.default(term=os.environ.get("TERM")), style=style_from_pygments_cls(get_style_by_name( settings.color_scheme)), editing_mode="VI" if settings.editing_mode in ["vim", "vi"] else "EMACS", history=history, enable_history_search=True, history_search_no_duplicates=settings.history_search_no_duplicates, search_ignore_case=settings.history_search_ignore_case, enable_suspend=True, tempfile_suffix=".R", input=CustomInput(sys.stdin), output=output, inputhook=get_inputhook(), mode_class=RadianMode) apply_settings(session, settings) def browse_activator(session): message = session.prompt_text if BROWSE_PATTERN.match(message): session.browse_level = BROWSE_PATTERN.match(message).group(1) return True else: return False def browse_on_pre_accept(session): if session.default_buffer.text.strip() in [ "n", "s", "f", "c", "cont", "Q", "where", "help" ]: session.add_history = False def shell_process_text(session): text = session.default_buffer.text shell.run_command(text) input_processors = [] if settings.highlight_matching_bracket: input_processors.append(HighlightMatchingBracketProcessor()) session.register_mode( "r", activator=lambda session: session.prompt_text == settings.prompt, insert_new_line=True, history_share_with="browse", get_message=lambda: settings.prompt, multiline=settings.indent_lines, complete_while_typing=settings.complete_while_typing, lexer=PygmentsLexer(SLexer), completer=RCompleter(timeout=settings.completion_timeout), key_bindings=create_key_bindings(), input_processors=input_processors, prompt_key_bindings=create_r_key_bindings(prase_text_complete)) session.register_mode("shell", on_post_accept=shell_process_text, insert_new_line=True, get_message=lambda: settings.shell_prompt, multiline=settings.indent_lines, complete_while_typing=settings.complete_while_typing, lexer=None, completer=SmartPathCompleter(), prompt_key_bindings=create_shell_key_bindings()) session.register_mode( "browse", activator=browse_activator, # on_pre_accept=browse_on_pre_accept, # disable insert_new_line=True, history_share_with="r", get_message=lambda: settings.browse_prompt.format(session.browse_level ), multiline=settings.indent_lines, complete_while_typing=True, lexer=PygmentsLexer(SLexer), completer=RCompleter(timeout=settings.completion_timeout), input_processors=input_processors, prompt_key_bindings=create_r_key_bindings(prase_text_complete), switchable_from=False, switchable_to=False) session.register_mode("unknown", insert_new_line=False, get_message=lambda: session.prompt_text, complete_while_typing=False, lexer=None, completer=None, prompt_key_bindings=None, switchable_from=False, switchable_to=False, input_processors=[]) return session
def run_cli(self): iterations = 0 sqlexecute = self.sqlexecute logger = self.logger self.configure_pager() if self.smart_completion: self.refresh_completions() author_file = os.path.join(PACKAGE_ROOT, 'AUTHORS') sponsor_file = os.path.join(PACKAGE_ROOT, 'SPONSORS') history_file = os.path.expanduser( os.environ.get('MYCLI_HISTFILE', '~/.mycli-history')) if dir_path_exists(history_file): history = FileHistory(history_file) else: history = None self.echo( 'Error: Unable to open the history file "{}". ' 'Your query history will not be saved.'.format(history_file), err=True, fg='red') key_bindings = mycli_bindings(self) if not self.less_chatty: print(' '.join(sqlexecute.server_type())) print('mycli', __version__) print('Chat: https://gitter.im/dbcli/mycli') print('Mail: https://groups.google.com/forum/#!forum/mycli-users') print('Home: http://mycli.net') print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file])) def get_message(): prompt = self.get_prompt(self.prompt_format) if self.prompt_format == self.default_prompt and len(prompt) > self.max_len_prompt: prompt = self.get_prompt('\\d> ') return [('class:prompt', prompt)] def get_continuation(width, line_number, is_soft_wrap): continuation = ' ' * (width - 1) + ' ' return [('class:continuation', continuation)] def show_suggestion_tip(): return iterations < 2 def one_iteration(text=None): if text is None: try: text = self.prompt_app.prompt() except KeyboardInterrupt: return special.set_expanded_output(False) try: text = self.handle_editor_command(text) except RuntimeError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') return if not text.strip(): return if self.destructive_warning: destroy = confirm_destructive_query(text) if destroy is None: pass # Query was not destructive. Nothing to do here. elif destroy is True: self.echo('Your call!') else: self.echo('Wise choice!') return # Keep track of whether or not the query is mutating. In case # of a multi-statement query, the overall query is considered # mutating if any one of the component statements is mutating mutating = False try: logger.debug('sql: %r', text) special.write_tee(self.get_prompt(self.prompt_format) + text) if self.logfile: self.logfile.write('\n# %s\n' % datetime.now()) self.logfile.write(text) self.logfile.write('\n') successful = False start = time() res = sqlexecute.run(text) self.formatter.query = text successful = True result_count = 0 for title, cur, headers, status in res: 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): self.echo('The result set has more than {} rows.'.format( threshold), fg='red') if not confirm('Do you want to continue?'): self.echo("Aborted!", err=True, fg='red') break if self.auto_vertical_output: max_width = self.prompt_app.output.get_size().columns else: max_width = None formatted = self.format_output( title, cur, headers, special.is_expanded_output(), max_width) t = time() - start try: if result_count > 0: self.echo('') try: self.output(formatted, status) except KeyboardInterrupt: pass if special.is_timing_enabled(): self.echo('Time: %0.03fs' % t) except KeyboardInterrupt: pass start = time() result_count += 1 mutating = mutating or is_mutating(status) special.unset_once_if_written() except EOFError as e: raise e except KeyboardInterrupt: # get last connection id connection_id_to_kill = sqlexecute.connection_id logger.debug("connection id to kill: %r", connection_id_to_kill) # Restart connection to the database sqlexecute.connect() try: for title, cur, headers, status in sqlexecute.run('kill %s' % connection_id_to_kill): status_str = str(status).lower() if status_str.find('ok') > -1: logger.debug("cancelled query, connection id: %r, sql: %r", connection_id_to_kill, text) self.echo("cancelled query", err=True, fg='red') except Exception as e: self.echo('Encountered error while cancelling query: {}'.format(e), err=True, fg='red') except NotImplementedError: self.echo('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.debug("Exception: %r", e) if (e.args[0] in (2003, 2006, 2013)): logger.debug('Attempting to reconnect.') self.echo('Reconnecting...', fg='yellow') try: sqlexecute.connect() logger.debug('Reconnected successfully.') one_iteration(text) return # OK to just return, cuz the recursion call runs to the end. except OperationalError as e: logger.debug('Reconnect failed. e: %r', e) self.echo(str(e), err=True, fg='red') # If reconnection failed, don't proceed further. return else: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') except Exception as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self.echo(str(e), err=True, fg='red') else: if is_dropping_database(text, self.sqlexecute.dbname): self.sqlexecute.dbname = None self.sqlexecute.connect() # Refresh the table names and column names if necessary. if need_completion_refresh(text): self.refresh_completions( reset=need_completion_reset(text)) finally: if self.logfile is False: self.echo("Warning: This query was not logged.", err=True, fg='red') query = Query(text, successful, mutating) self.query_history.append(query) get_toolbar_tokens = create_toolbar_tokens_func( self, show_suggestion_tip) if self.wider_completion_menu: complete_style = CompleteStyle.MULTI_COLUMN else: complete_style = CompleteStyle.COLUMN with self._completer_lock: if self.key_bindings == 'vi': editing_mode = EditingMode.VI else: editing_mode = EditingMode.EMACS self.prompt_app = PromptSession( lexer=PygmentsLexer(MyCliLexer), reserve_space_for_menu=self.get_reserved_space(), message=get_message, prompt_continuation=get_continuation, bottom_toolbar=get_toolbar_tokens, complete_style=complete_style, input_processors=[ConditionalProcessor( processor=HighlightMatchingBracketProcessor( chars='[](){}'), filter=HasFocus(DEFAULT_BUFFER) & ~IsDone() )], tempfile_suffix='.sql', completer=DynamicCompleter(lambda: self.completer), history=history, auto_suggest=AutoSuggestFromHistory(), complete_while_typing=True, multiline=cli_is_multiline(self), style=style_factory(self.syntax_style, self.cli_style), include_default_pygments_style=False, key_bindings=key_bindings, enable_open_in_editor=True, enable_system_prompt=True, enable_suspend=True, editing_mode=editing_mode, search_ignore_case=True ) try: while True: one_iteration() iterations += 1 except EOFError: special.close_tee() if not self.less_chatty: self.echo('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