def run_schema_metadata_validation(self, charset, test_db): try: # setup db test_db_name = test_db client = self.create_mssqlcliclient(test_db_name) new_schemas = [] for characters in self.get_next_characters( self.charset_dict[charset]): schema_name = u'mssqlcli_{0}_{1}_{2}'.format( self.get_local_machine_name(), random_str(), characters) query = u'CREATE SCHEMA {0}'.format(schema_name) if self.run_query(client, query): new_schemas.append(schema_name) else: assert False # should not fail completer = MssqlCompleter(smart_completion=True) completion_refresher.refresh_schemas(completer, client) completions = completer.get_completions(document=Document( u'select * from ', 14, None), complete_event=None, smart_completion=True) db_schemas = set(map(lambda e: e.text, completions)) for new_schema in new_schemas: assert u'"{}"'.format(new_schema) in db_schemas finally: shutdown(client)
def get_completer(self, settings=None, casing=None): # pylint: disable=too-many-locals metadata = self.metadata comp = MssqlCompleter(smart_completion=True, settings=settings) schemas, tables, tbl_cols, views, view_cols = [], [], [], [], [] for sch, tbls in metadata['tables'].items(): schemas.append(sch) for tbl, cols in tbls.items(): tables.append((sch, tbl)) # Let all columns be text columns tbl_cols.extend( [self._make_col(sch, tbl, col) for col in cols]) for sch, tbls in metadata.get('views', {}).items(): for tbl, cols in tbls.items(): views.append((sch, tbl)) # Let all columns be text columns view_cols.extend( [self._make_col(sch, tbl, col) for col in cols]) functions = [ FunctionMetadata(sch, *func_meta, arg_defaults=None) for sch, funcs in metadata['functions'].items() for func_meta in funcs ] datatypes = [(sch, typ) for sch, datatypes in metadata['datatypes'].items() for typ in datatypes] foreignkeys = [ ForeignKey(*fk) for fks in metadata['foreignkeys'].values() for fk in fks ] comp.extend_schemas(schemas) comp.extend_relations(tables, kind='tables') comp.extend_relations(views, kind='views') comp.extend_columns(tbl_cols, kind='tables') comp.extend_columns(view_cols, kind='views') comp.extend_functions(functions) comp.extend_datatypes(datatypes) comp.extend_foreignkeys(foreignkeys) comp.set_search_path(['public']) comp.extend_casing(casing or []) return comp
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
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 __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) 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.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.")
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