def run(conn, sql, config, user_namespace): if sql.strip(): for statement in sqlparse.split(sql): first_word = sql.strip().split()[0].lower() if first_word == 'begin': raise Exception("ipython_sql does not support transactions") if first_word.startswith('\\') and 'postgres' in str(conn.dialect): pgspecial = PGSpecial() _, cur, headers, _ = pgspecial.execute( conn.session.connection.cursor(), statement)[0] result = FakeResultProxy(cur, headers) else: txt = sqlalchemy.sql.text(statement) result = conn.session.execute(txt, user_namespace) try: # mssql has autocommit if config.autocommit and ('mssql' not in str(conn.dialect)): conn.session.execute('commit') except sqlalchemy.exc.OperationalError: pass # not all engines can commit if result and config.feedback: print(interpret_rowcount(result.rowcount)) resultset = ResultSet(result, statement, config) if config.autopandas: return resultset.DataFrame() else: return resultset #returning only last result, intentionally else: return 'Connected: %s' % conn.name
def test_exit_without_active_connection(executor): quit_handler = MagicMock() pgspecial = PGSpecial() pgspecial.register(quit_handler, '\\q', '\\q', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=True, aliases=(':q',)) with patch.object(executor, "conn", BrokenConnection()): # we should be able to quit the app, even without active connection run(executor, "\\q", pgspecial=pgspecial) quit_handler.assert_called_once() # an exception should be raised when running a query without active connection with pytest.raises(psycopg2.InterfaceError): run(executor, "select 1", pgspecial=pgspecial)
def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute # Load config. c = self.config = get_config(pgclirc_file) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.table_format = c['main']['table_format'] self.syntax_style = c['main']['syntax_style'] self.cli_style = c['colors'] self.wider_completion_menu = c['main'].as_bool('wider_completion_menu') self.less_chatty = c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = c['main'].get('prompt', self.default_prompt) self.on_error = c['main']['on_error'].upper() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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']} completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None
def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) pgclirc_file = pgclirc_file or '%sconfig' % config_location() default_config = os.path.join(package_root, 'pgclirc') write_default_config(default_config, pgclirc_file) # Load config. c = self.config = load_config(pgclirc_file, default_config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') 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.on_error = c['main']['on_error'].upper() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None
def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) default_config = os.path.join(package_root, "pgclirc") write_default_config(default_config, pgclirc_file) self.pgspecial = PGSpecial() # Load config. c = self.config = load_config(pgclirc_file, default_config) self.multi_line = c["main"].as_bool("multi_line") self.vi_mode = c["main"].as_bool("vi") self.pgspecial.timing_enabled = c["main"].as_bool("timing") 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.on_error = c["main"]["on_error"].upper() self.completion_refresher = CompletionRefresher() self.logger = logging.getLogger(__name__) self.initialize_logging() self.query_history = [] # Initialize completer smart_completion = c["main"].as_bool("smart_completion") completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.cli = None
def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None # Load config. c = self.config = get_config(pgclirc_file) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.multiline_mode = c['main'].get('multi_line_mode', 'psql') self.vi_mode = c['main'].as_bool('vi') self.auto_expand = auto_vertical_output or c['main'].as_bool( 'auto_expand') self.expanded_output = c['main'].as_bool('expand') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.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') c_dest_warning = c['main'].as_bool('destructive_warning') self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt) self.prompt_dsn_format = prompt_dsn self.on_error = c['main']['on_error'].upper() self.decimal_format = c['data_formats']['decimal'] self.float_format = c['data_formats']['float'] self.pgspecial.pset_pager(self.config['main'].as_bool( 'enable_pager') and "on" or "off") self.style_output = style_factory_output( self.syntax_style, c['colors']) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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': single_connection, 'less_chatty': less_chatty, 'keyword_casing': keyword_casing, } completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None
class PGCli(object): def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) default_config = os.path.join(package_root, 'pgclirc') write_default_config(default_config, pgclirc_file) self.pgspecial = PGSpecial() # Load config. c = self.config = load_config(pgclirc_file, default_config) self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') 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.on_error = c['main']['on_error'].upper() self.completion_refresher = CompletionRefresher() self.logger = logging.getLogger(__name__) self.initialize_logging() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.cli = None def register_special_commands(self): self.pgspecial.register(self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) self.pgspecial.register(self.refresh_completions, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.refresh_completions, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') def change_db(self, pattern, **_): if pattern: db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern self.pgexecute.connect(database=db) else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e))] return self.pgexecute.run(query, self.pgspecial, on_error=self.on_error) def initialize_logging(self): log_file = self.config['main']['log_file'] log_level = self.config['main']['log_level'] level_map = {'CRITICAL': logging.CRITICAL, 'ERROR': logging.ERROR, 'WARNING': logging.WARNING, 'INFO': logging.INFO, 'DEBUG': logging.DEBUG } handler = logging.FileHandler(os.path.expanduser(log_file)) formatter = logging.Formatter( '%(asctime)s (%(process)d/%(threadName)s) ' '%(name)s %(levelname)s - %(message)s') handler.setFormatter(formatter) root_logger = logging.getLogger('pgcli') root_logger.addHandler(handler) root_logger.setLevel(level_map[log_level.upper()]) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash self.connect(database, uri.hostname, uri.username, uri.port, uri.password) def connect(self, database='', host='', user='', port='', passwd='', dsn=''): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn) except OperationalError as e: if ('no password supplied' in utf8tounicode(e.args[0]) and auto_passwd_prompt): passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute def handle_editor_command(self, cli, document): """ 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 """ while special.editor_command(document.text): filename = special.get_filename(document.text) sql, message = special.open_external_editor(filename, sql=document.text) if message: # Something went wrong. Raise an exception and bail. raise RuntimeError(message) cli.current_buffer.document = Document(sql, cursor_position=len(sql)) document = cli.run(False) continue return document def run_cli(self): logger = self.logger original_less_opts = self.adjust_less_opts() self.refresh_completions() self.cli = self._build_cli() print('Version:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') try: while True: document = self.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(self.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 # Initialize default metaquery in case execution fails query = MetaQuery(query=document.text, successful=False) try: output, query = self._evaluate_command(document.text) except KeyboardInterrupt: # Restart connection to the database self.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: if ('server closed the connection' in utf8tounicode(e.args[0])): self._handle_server_closed_connection() 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: try: click.echo_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: 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: self.refresh_completions(reset=True) elif query.meta_changed: self.refresh_completions(reset=False) elif query.path_changed: logger.debug('Refreshing search path') with self._completer_lock: self.completer.set_search_path( self.pgexecute.search_path()) logger.debug('Search path: %r', self.completer.search_path) 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): 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(_): return [(Token.Prompt, '%s> ' % self.pgexecute.dbname)] get_toolbar_tokens = create_toolbar_tokens_func( lambda: self.vi_mode, self.completion_refresher.is_refreshing) layout = create_default_layout( lexer=PostgresLexer, reserve_space_for_menu=True, get_prompt_tokens=prompt_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()), ]) history_file = self.config['main']['history_file'] with self._completer_lock: buf = PGBuffer( always_multiline=self.multi_line, completer=self.completer, history=FileHistory(os.path.expanduser(history_file)), 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()) return cli 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) """ logger = self.logger logger.debug('sql: %r', text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 # Run the query. start = time() on_error_resume = self.on_error == 'RESUME' res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success 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): 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 if self.pgspecial.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output( title, cur, headers, status, self.table_format, self.pgspecial.expanded_output, max_width) output.extend(formatted) end = time() total += end - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, meta_changed, db_changed, path_changed, mutated) return output, meta_query 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.pgexecute.connect() click.secho('Reconnected!\nTry the command again.', fg='green') except OperationalError as e: click.secho(str(e), err=True, fg='red') def adjust_less_opts(self): less_opts = os.environ.get('LESS', '') self.logger.debug('Original value for LESS env var: %r', less_opts) os.environ['LESS'] = '-SRXF' return less_opts def refresh_completions(self, reset=False): if reset: with self._completer_lock: self.completer.reset_completions() self.completion_refresher.refresh(self.pgexecute, self.pgspecial, self._on_completions_refreshed) return [(None, None, None, 'Auto-completion refresh started in the background.')] def _on_completions_refreshed(self, new_completer): self._swap_completer_objects(new_completer) 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): """Swap the completer object in cli with the newly created completer. """ with self._completer_lock: self.completer = new_completer # When pgcli 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 __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None self.watch_command = None # Load config. c = self.config = get_config(pgclirc_file) NamedQueries.instance = NamedQueries.from_config(self.config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.multiline_mode = c['main'].get('multi_line_mode', 'psql') self.vi_mode = c['main'].as_bool('vi') self.auto_expand = auto_vertical_output or c['main'].as_bool( 'auto_expand') self.expanded_output = c['main'].as_bool('expand') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.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') c_dest_warning = c['main'].as_bool('destructive_warning') self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt) self.prompt_dsn_format = prompt_dsn self.on_error = c['main']['on_error'].upper() self.decimal_format = c['data_formats']['decimal'] self.float_format = c['data_formats']['float'] self.initialize_keyring() self.pgspecial.pset_pager(self.config['main'].as_bool( 'enable_pager') and "on" or "off") self.style_output = style_factory_output( self.syntax_style, c['colors']) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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': single_connection, 'less_chatty': less_chatty, 'keyword_casing': keyword_casing, } completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.prompt_app = None
class PGCli(object): default_prompt = '\\u@\\h:\\d> ' max_len_prompt = 30 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, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None self.watch_command = None # Load config. c = self.config = get_config(pgclirc_file) NamedQueries.instance = NamedQueries.from_config(self.config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.multiline_mode = c['main'].get('multi_line_mode', 'psql') self.vi_mode = c['main'].as_bool('vi') self.auto_expand = auto_vertical_output or c['main'].as_bool( 'auto_expand') self.expanded_output = c['main'].as_bool('expand') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.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') c_dest_warning = c['main'].as_bool('destructive_warning') self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt) self.prompt_dsn_format = prompt_dsn self.on_error = c['main']['on_error'].upper() self.decimal_format = c['data_formats']['decimal'] self.float_format = c['data_formats']['float'] self.initialize_keyring() self.pgspecial.pset_pager(self.config['main'].as_bool( 'enable_pager') and "on" or "off") self.style_output = style_factory_output( self.syntax_style, c['colors']) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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': single_connection, 'less_chatty': less_chatty, 'keyword_casing': keyword_casing, } completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.prompt_app = None def quit(self): raise PgCliQuitError def register_special_commands(self): self.pgspecial.register( self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) refresh_callback = lambda: self.refresh_completions( persist_priorities='all') self.pgspecial.register(self.quit, '\\q', '\\q', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=True, aliases=(':q',)) self.pgspecial.register(self.quit, 'quit', 'quit', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=False, aliases=('exit',)) self.pgspecial.register(refresh_callback, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(refresh_callback, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') self.pgspecial.register(self.write_to_file, '\\o', '\\o [filename]', 'Send all query results to file.') self.pgspecial.register(self.info_connection, '\\conninfo', '\\conninfo', 'Get connection details') self.pgspecial.register(self.change_table_format, '\\T', '\\T [format]', 'Change the table format used to output results') def change_table_format(self, pattern, **_): try: if pattern not in TabularOutputFormatter().supported_formats: raise ValueError() self.table_format = pattern yield (None, None, None, 'Changed table format to {}'.format(pattern)) except ValueError: msg = 'Table format {} not recognized. Allowed formats:'.format( pattern) for table_type in TabularOutputFormatter().supported_formats: msg += "\n\t{}".format(table_type) msg += '\nCurrently set to: %s' % self.table_format yield (None, None, None, msg) def info_connection(self, **_): if self.pgexecute.host.startswith('/'): host = 'socket "%s"' % self.pgexecute.host else: host = 'host "%s"' % self.pgexecute.host yield (None, None, None, 'You are connected to database "%s" as user ' '"%s" on %s at port "%s".' % (self.pgexecute.dbname, self.pgexecute.user, host, self.pgexecute.port)) def change_db(self, pattern, **_): if pattern: # Get all the parameters in pattern, handling double quotes if any. infos = re.findall(r'"[^"]*"|[^"\'\s]+', pattern) # Now removing quotes. list(map(lambda s: s.strip('"'), infos)) infos.extend([None] * (4 - len(infos))) db, user, host, port = infos try: self.pgexecute.connect(database=db, user=user, host=host, port=port, **self.pgexecute.extra_args) except OperationalError as e: click.secho(str(e), err=True, fg='red') click.echo("Previous connection kept") else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message, '', False, True)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e), '', False, True)] if (self.destructive_warning and confirm_destructive_query(query) is False): message = 'Wise choice. Command execution stopped.' return [(None, None, None, message)] on_error_resume = (self.on_error == 'RESUME') return self.pgexecute.run( query, self.pgspecial, on_error_resume=on_error_resume ) def write_to_file(self, pattern, **_): if not pattern: self.output_file = None message = 'File output disabled' return [(None, None, None, message, '', True, 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, True)] self.output_file = filename message = 'Writing to file "%s"' % self.output_file return [(None, None, None, message, '', True, True)] def initialize_logging(self): log_file = self.config['main']['log_file'] if log_file == 'default': log_file = config_location() + '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('pgcli') root_logger.addHandler(handler) root_logger.setLevel(log_level) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) pgspecial_logger = logging.getLogger('pgspecial') pgspecial_logger.addHandler(handler) pgspecial_logger.setLevel(log_level) def initialize_keyring(self): global keyring keyring_enabled = self.config["main"].as_bool("keyring") if keyring_enabled: # Try best to load keyring (issue #1041). import importlib try: keyring = importlib.import_module('keyring') except Exception as e: # ImportError for Python 2, ModuleNotFoundError for Python 3 self.logger.warning('import keyring failed: %r.', e) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): kwargs = psycopg2.extensions.parse_dsn(uri) remap = { 'dbname': 'database', 'password': '******', } kwargs = {remap.get(k, k): v for k, v in kwargs.items()} self.connect(**kwargs) def connect(self, database='', host='', user='', port='', passwd='', dsn='', **kwargs): # Connect to the database. if not user: user = getuser() if not database: database = user kwargs.setdefault('application_name', 'pgcli') # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Find password from store key = '%s@%s' % (user, host) keyring_error_message = dedent("""\ {} {} To remove this message do one of the following: - prepare keyring as described at: https://keyring.readthedocs.io/en/stable/ - uninstall keyring: pip uninstall keyring - disable keyring in our configuration: add keyring = False to [main]""") if not passwd and keyring: try: passwd = keyring.get_password('pgcli', key) except ( RuntimeError, keyring.errors.InitError ) as e: click.secho( keyring_error_message.format( "Load your password from keyring returned:", str(e) ), err=True, fg='red' ) # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password for %s' % user, hide_input=True, show_default=False, type=str) def should_ask_for_password(exc): # Prompt for a password after 1st attempt to connect # fails. Don't prompt if the -w flag is supplied if self.never_passwd_prompt: return False error_msg = utf8tounicode(exc.args[0]) if "no password supplied" in error_msg: return True if "password authentication failed" in error_msg: return True return False # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing or incorrect password, but we're allowed to # prompt for a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn, **kwargs) except (OperationalError, InterfaceError) as e: if should_ask_for_password(e): passwd = click.prompt('Password for %s' % user, hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn, **kwargs) else: raise e if passwd and keyring: try: keyring.set_password('pgcli', key, passwd) except ( RuntimeError, keyring.errors.KeyringError, ) as e: click.secho( keyring_error_message.format( "Set password in keyring returned:", str(e) ), err=True, fg='red' ) except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute 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 """ editor_command = special.editor_command(text) while editor_command: if editor_command == '\\e': filename = special.get_filename(text) query = special.get_editor_query( text) or self.get_last_query() else: # \ev or \ef filename = None spec = text.split()[1] if editor_command == '\\ev': query = self.pgexecute.view_definition(spec) elif editor_command == '\\ef': query = self.pgexecute.function_definition(spec) 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_app.prompt(default=sql) break except KeyboardInterrupt: sql = "" editor_command = special.editor_command(text) return text def execute_command(self, text): logger = self.logger query = MetaQuery(query=text, successful=False) try: if (self.destructive_warning): destroy = confirm = confirm_destructive_query(text) if destroy is False: click.secho('Wise choice!') raise KeyboardInterrupt elif destroy: click.secho('Your call!') output, query = self._evaluate_command(text) except KeyboardInterrupt: # Restart connection to the database self.pgexecute.connect() logger.debug("cancelled query, sql: %r", text) click.secho("cancelled query", err=True, fg='red') except NotImplementedError: click.secho('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self._handle_server_closed_connection(text) except (PgCliQuitError, EOFError) as e: raise 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: self.echo_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second if query.total_time > 1: print('Time: %0.03fs (%s), executed in: %0.03fs (%s)' % (query.total_time, humanize.time.naturaldelta( query.total_time), query.execution_time, humanize.time.naturaldelta(query.execution_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') elif query.path_changed: logger.debug('Refreshing search path') with self._completer_lock: self.completer.set_search_path( self.pgexecute.search_path()) logger.debug('Search path: %r', self.completer.search_path) return query def run_cli(self): logger = self.logger history_file = self.config['main']['history_file'] if history_file == 'default': history_file = config_location() + 'history' history = FileHistory(os.path.expanduser(history_file)) self.refresh_completions(history=history, persist_priorities='none') self.prompt_app = self._build_cli(history) if not self.less_chatty: print('Server: PostgreSQL', self.pgexecute.server_version) print('Version:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') try: while True: try: text = self.prompt_app.prompt() except KeyboardInterrupt: continue 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()) click.secho(str(e), err=True, fg='red') continue # Initialize default metaquery in case execution fails self.watch_command, timing = special.get_watch_command(text) if self.watch_command: while self.watch_command: try: query = self.execute_command(self.watch_command) click.echo( 'Waiting for {0} seconds before repeating' .format(timing)) sleep(timing) except KeyboardInterrupt: self.watch_command = None else: query = self.execute_command(text) self.now = dt.datetime.today() # Allow PGCompleter to learn user's preferred keywords, etc. with self._completer_lock: self.completer.extend_query_history(text) self.query_history.append(query) except (PgCliQuitError, EOFError): if not self.less_chatty: print ('Goodbye!') 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', multiline=pg_is_multiline(self), 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 _should_show_limit_prompt(self, status, cur): """returns True if limit prompt should be shown, False otherwise.""" if not is_select(status): return False return self.row_limit > 0 and cur and (cur.rowcount > self.row_limit) def _evaluate_command(self, text): """Used to run a command entered by the user during CLI operation (Puts the E in REPL) returns (results, MetaQuery) """ logger = self.logger logger.debug('sql: %r', text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 execution = 0 # Run the query. start = time() on_error_resume = self.on_error == 'RESUME' res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success, is_special in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) threshold = self.row_limit if self._should_show_limit_prompt(status, cur): click.secho('The result set has more than %s rows.' % threshold, fg='red') if not click.confirm('Do you want to continue?'): click.secho("Aborted!", err=True, fg='red') break if self.pgspecial.auto_expand or self.auto_expand: max_width = self.prompt_app.output.get_size().columns else: max_width = None expanded = self.pgspecial.expanded_output or self.expanded_output settings = OutputSettings( table_format=self.table_format, dcmlfmt=self.decimal_format, floatfmt=self.float_format, missingval=self.null_string, expanded=expanded, max_width=max_width, case_function=( self.completer.case if self.settings['case_column_headers'] else lambda x: x ), style_output=self.style_output ) execution = time() - start formatted = format_output(title, cur, headers, status, settings) output.extend(formatted) total = time() - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, execution, meta_changed, db_changed, path_changed, mutated, is_special) return output, meta_query def _handle_server_closed_connection(self, text): """Used during CLI execution.""" try: click.secho('Reconnecting...', fg='green') self.pgexecute.connect() click.secho('Reconnected!', fg='green') self.execute_command(text) except OperationalError as e: click.secho('Reconnect Failed', fg='red') click.secho(str(e), err=True, fg='red') def refresh_completions(self, history=None, persist_priorities='all'): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial(self._on_completions_refreshed, persist_priorities=persist_priorities) self.completion_refresher.refresh(self.pgexecute, self.pgspecial, 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_app: # After refreshing, redraw the CLI to clear the statusbar # "Refreshing completions..." indicator self.prompt_app.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 self.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): # should be before replacing \\d string = string.replace('\\dsn_alias', self.dsn_alias or '') string = string.replace('\\t', self.now.strftime('%x %X')) string = string.replace('\\u', self.pgexecute.user or '(none)') string = string.replace('\\H', self.pgexecute.host or '(none)') string = string.replace('\\h', self.pgexecute.short_host or '(none)') string = string.replace('\\d', self.pgexecute.dbname or '(none)') string = string.replace('\\p', str( self.pgexecute.port) if self.pgexecute.port is not None else '5432') string = string.replace('\\i', str(self.pgexecute.pid) or '(none)') string = string.replace('\\#', "#" if (self.pgexecute.superuser) else ">") 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 is_too_wide(self, line): """Will this line be too wide to fit into terminal?""" if not self.prompt_app: return False return len(COLOR_CODE_REGEX.sub('', line)) > self.prompt_app.output.get_size().columns def is_too_tall(self, lines): """Are there too many lines to fit into terminal?""" if not self.prompt_app: return False return len(lines) >= (self.prompt_app.output.get_size().rows - 4) def echo_via_pager(self, text, color=None): if self.pgspecial.pager_config == PAGER_OFF or self.watch_command: click.echo(text, color=color) elif self.pgspecial.pager_config == PAGER_LONG_OUTPUT: lines = text.split('\n') # The last 4 lines are reserved for the pgcli menu and padding if self.is_too_tall(lines) or any(self.is_too_wide(l) for l in lines): click.echo_via_pager(text, color=color) else: click.echo(text, color=color) else: click.echo_via_pager(text, color)
class PGCli(object): 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: ' + '\'' + configured_pager + '\'') os.environ['PAGER'] = configured_pager elif os_environ_pager: self.logger.info('Default pager found in PAGER environment variable: ' + '\'' + os_environ_pager + '\'') os.environ['PAGER'] = os_environ_pager else: self.logger.info('No default pager found in environment. Using os default pager') # Always set default set of less recommended options, they are ignored if pager is # different than less or is already parameterized with their own arguments os.environ['LESS'] = '-SRXF' def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute # Load config. c = self.config = get_config(pgclirc_file) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.table_format = c['main']['table_format'] self.syntax_style = c['main']['syntax_style'] self.cli_style = c['colors'] self.wider_completion_menu = c['main'].as_bool('wider_completion_menu') self.less_chatty = c['main'].as_bool('less_chatty') self.on_error = c['main']['on_error'].upper() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') settings = {'casing_file': get_casing_file(c), 'generate_casing_file': c['main'].as_bool('generate_casing_file'), 'asterisk_column_order': c['main']['asterisk_column_order']} completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None def register_special_commands(self): self.pgspecial.register( self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) refresh_callback = lambda: self.refresh_completions( persist_priorities='all') self.pgspecial.register(refresh_callback, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(refresh_callback, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') self.pgspecial.register(self.write_to_file, '\\o', '\\o [filename]', 'Send all query results to file.') def change_db(self, pattern, **_): if pattern: db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern self.pgexecute.connect(database=db) else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message, '', False)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e), '', False)] on_error_resume = (self.on_error == 'RESUME') return self.pgexecute.run( query, self.pgspecial, on_error_resume=on_error_resume ) 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() + '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 = 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('pgcli') root_logger.addHandler(handler) root_logger.setLevel(log_level) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) pgspecial_logger = logging.getLogger('pgspecial') pgspecial_logger.addHandler(handler) pgspecial_logger.setLevel(log_level) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash arguments = [database, uri.hostname, uri.username, uri.port, uri.password] # unquote each URI part (they may be percent encoded) self.connect(*list(map(lambda p: unquote(p) if p else p, arguments))) def connect(self, database='', host='', user='', port='', passwd='', dsn=''): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn) except OperationalError as e: if ('no password supplied' in utf8tounicode(e.args[0]) and auto_passwd_prompt): passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute def handle_editor_command(self, cli, document): """ 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 """ while special.editor_command(document.text): filename = special.get_filename(document.text) sql, message = special.open_external_editor(filename, sql=document.text) if message: # Something went wrong. Raise an exception and bail. raise RuntimeError(message) cli.current_buffer.document = Document(sql, cursor_position=len(sql)) document = cli.run(False) continue return document def run_cli(self): logger = self.logger history_file = self.config['main']['history_file'] if history_file == 'default': history_file = config_location() + 'history' history = FileHistory(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:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') try: while True: document = self.cli.run(True) # 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(self.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 # Initialize default metaquery in case execution fails query = MetaQuery(query=document.text, successful=False) try: output, query = self._evaluate_command(document.text) except KeyboardInterrupt: # Restart connection to the database self.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: if ('server closed the connection' in utf8tounicode(e.args[0])): self._handle_server_closed_connection() 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: try: if self.output_file and not document.text.startswith(('\\o ', '\\? ')): try: with open(self.output_file, 'a', encoding='utf-8') as f: click.echo(document.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 self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second 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') elif query.path_changed: logger.debug('Refreshing search path') with self._completer_lock: self.completer.set_search_path( self.pgexecute.search_path()) logger.debug('Search path: %r', self.completer.search_path) # Allow PGCompleter 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: if not self.less_chatty: print ('Goodbye!') 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(_): return [(Token.Prompt, '%s> ' % self.pgexecute.dbname)] def get_continuation_tokens(cli, width): return [(Token.Continuation, '.' * (width - 1) + ' ')] get_toolbar_tokens = create_toolbar_tokens_func( lambda: self.vi_mode, self.completion_refresher.is_refreshing) layout = create_prompt_layout( lexer=PygmentsLexer(PostgresLexer), reserve_space_for_menu=4, 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( always_multiline=self.multi_line, 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, cur): """returns True if limit prompt should be shown, False otherwise.""" if not is_select(status): return False return self.row_limit > 0 and cur and cur.rowcount > self.row_limit def _evaluate_command(self, text): """Used to run a command entered by the user during CLI operation (Puts the E in REPL) returns (results, MetaQuery) """ logger = self.logger logger.debug('sql: %r', text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 # Run the query. start = time() on_error_resume = self.on_error == 'RESUME' res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) threshold = self.row_limit if self._should_show_limit_prompt(status, cur): click.secho('The result set has more than %s rows.' % threshold, fg='red') if not click.confirm('Do you want to continue?'): click.secho("Aborted!", err=True, fg='red') break if self.pgspecial.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output( title, cur, headers, status, self.table_format, self.pgspecial.expanded_output, max_width) output.extend(formatted) end = time() total += end - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, meta_changed, db_changed, path_changed, mutated) return output, meta_query 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.pgexecute.connect() click.secho('Reconnected!\nTry the command again.', fg='green') except OperationalError as e: click.secho(str(e), err=True, fg='red') def refresh_completions(self, history=None, persist_priorities='all'): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial(self._on_completions_refreshed, persist_priorities=persist_priorities) c = self.config settings = {'casing_file': get_casing_file(c), 'generate_casing_file': c['main'].as_bool('generate_casing_file'), 'asterisk_column_order': c['main']['asterisk_column_order']} self.completion_refresher.refresh(self.pgexecute, self.pgspecial, callback, history=history, settings=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 pgcli 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)
class PGCli(object): def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) default_config = os.path.join(package_root, "pgclirc") write_default_config(default_config, pgclirc_file) self.pgspecial = PGSpecial() # Load config. c = self.config = load_config(pgclirc_file, default_config) self.multi_line = c["main"].as_bool("multi_line") self.vi_mode = c["main"].as_bool("vi") self.pgspecial.timing_enabled = c["main"].as_bool("timing") 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.on_error = c["main"]["on_error"].upper() self.completion_refresher = CompletionRefresher() self.logger = logging.getLogger(__name__) self.initialize_logging() self.query_history = [] # Initialize completer smart_completion = c["main"].as_bool("smart_completion") completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.cli = None def register_special_commands(self): self.pgspecial.register( self.change_db, "\\c", "\\c[onnect] database_name", "Change to a new database.", aliases=("use", "\\connect", "USE"), ) refresh_callback = lambda: self.refresh_completions(persist_priorities="all") self.pgspecial.register(refresh_callback, "\\#", "\\#", "Refresh auto-completions.", arg_type=NO_QUERY) self.pgspecial.register( refresh_callback, "\\refresh", "\\refresh", "Refresh auto-completions.", arg_type=NO_QUERY ) self.pgspecial.register(self.execute_from_file, "\\i", "\\i filename", "Execute commands from file.") def change_db(self, pattern, **_): if pattern: db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern self.pgexecute.connect(database=db) else: self.pgexecute.connect() yield ( None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user), ) def execute_from_file(self, pattern, **_): if not pattern: message = "\\i: missing required argument" return [(None, None, None, message)] try: with open(os.path.expanduser(pattern), encoding="utf-8") as f: query = f.read() except IOError as e: return [(None, None, None, str(e))] return self.pgexecute.run(query, self.pgspecial, on_error=self.on_error) def initialize_logging(self): log_file = self.config["main"]["log_file"] if log_file == "default": log_file = config_location() + "log" log_level = self.config["main"]["log_level"] level_map = { "CRITICAL": logging.CRITICAL, "ERROR": logging.ERROR, "WARNING": logging.WARNING, "INFO": logging.INFO, "DEBUG": logging.DEBUG, } handler = logging.FileHandler(os.path.expanduser(log_file)) formatter = logging.Formatter( "%(asctime)s (%(process)d/%(threadName)s) " "%(name)s %(levelname)s - %(message)s" ) handler.setFormatter(formatter) root_logger = logging.getLogger("pgcli") root_logger.addHandler(handler) root_logger.setLevel(level_map[log_level.upper()]) root_logger.debug("Initializing pgcli logging.") root_logger.debug("Log file %r.", log_file) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash self.connect(database, uri.hostname, uri.username, uri.port, uri.password) def connect(self, database="", host="", user="", port="", passwd="", dsn=""): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get("PGPASSWORD", "") # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt("Password", hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn) except OperationalError as e: if "no password supplied" in utf8tounicode(e.args[0]) and auto_passwd_prompt: passwd = click.prompt("Password", hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug("Database connection failed: %r.", e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg="red") exit(1) self.pgexecute = pgexecute def handle_editor_command(self, cli, document): """ 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 """ while special.editor_command(document.text): filename = special.get_filename(document.text) sql, message = special.open_external_editor(filename, sql=document.text) if message: # Something went wrong. Raise an exception and bail. raise RuntimeError(message) cli.current_buffer.document = Document(sql, cursor_position=len(sql)) document = cli.run(False) continue return document def run_cli(self): logger = self.logger original_less_opts = self.adjust_less_opts() history_file = self.config["main"]["history_file"] if history_file == "default": history_file = config_location() + "history" history = FileHistory(os.path.expanduser(history_file)) self.refresh_completions(history=history, persist_priorities="none") self.cli = self._build_cli(history) print("Version:", __version__) print("Chat: https://gitter.im/dbcli/pgcli") print("Mail: https://groups.google.com/forum/#!forum/pgcli") print("Home: http://pgcli.com") try: while True: document = self.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(self.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 # Initialize default metaquery in case execution fails query = MetaQuery(query=document.text, successful=False) try: output, query = self._evaluate_command(document.text) except KeyboardInterrupt: # Restart connection to the database self.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: if "server closed the connection" in utf8tounicode(e.args[0]): self._handle_server_closed_connection() 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: try: click.echo_via_pager("\n".join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second 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") elif query.path_changed: logger.debug("Refreshing search path") with self._completer_lock: self.completer.set_search_path(self.pgexecute.search_path()) logger.debug("Search path: %r", self.completer.search_path) # Allow PGCompleter 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: 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(_): return [(Token.Prompt, "%s> " % self.pgexecute.dbname)] get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.vi_mode, self.completion_refresher.is_refreshing) layout = create_prompt_layout( lexer=PygmentsLexer(PostgresLexer), reserve_space_for_menu=4, get_prompt_tokens=prompt_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( always_multiline=self.multi_line, completer=self.completer, history=history, complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT, ) 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, ) cli = CommandLineInterface(application=application) return cli 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) """ logger = self.logger logger.debug("sql: %r", text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 # Run the query. start = time() on_error_resume = self.on_error == "RESUME" res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success 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: 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 if self.pgspecial.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output( title, cur, headers, status, self.table_format, self.pgspecial.expanded_output, max_width ) output.extend(formatted) end = time() total += end - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, meta_changed, db_changed, path_changed, mutated) return output, meta_query 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.pgexecute.connect() click.secho("Reconnected!\nTry the command again.", fg="green") except OperationalError as e: click.secho(str(e), err=True, fg="red") def adjust_less_opts(self): less_opts = os.environ.get("LESS", "") self.logger.debug("Original value for LESS env var: %r", less_opts) os.environ["LESS"] = "-SRXF" return less_opts def refresh_completions(self, history=None, persist_priorities="all"): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial(self._on_completions_refreshed, persist_priorities=persist_priorities) self.completion_refresher.refresh(self.pgexecute, self.pgspecial, callback, history=history) 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 pgcli 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 __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, prompt=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute # Load config. c = self.config = get_config(pgclirc_file) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.multiline_mode = c['main'].get('multi_line_mode', 'psql') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.table_format = c['main']['table_format'] self.syntax_style = c['main']['syntax_style'] self.cli_style = c['colors'] self.wider_completion_menu = c['main'].as_bool('wider_completion_menu') self.less_chatty = c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = prompt if prompt is not None else c['main'].get( 'prompt', self.default_prompt) self.on_error = c['main']['on_error'].upper() self.decimal_format = c['data_formats']['decimal'] self.float_format = c['data_formats']['float'] self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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'], 'single_connection': single_connection, 'keyword_casing': keyword_casing, } completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None
class PGCli(object): default_prompt = '\\u@\\h:\\d> ' max_len_prompt = 30 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, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None # Load config. c = self.config = get_config(pgclirc_file) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.multiline_mode = c['main'].get('multi_line_mode', 'psql') self.vi_mode = c['main'].as_bool('vi') self.auto_expand = auto_vertical_output or c['main'].as_bool( 'auto_expand') self.expanded_output = c['main'].as_bool('expand') self.pgspecial.timing_enabled = c['main'].as_bool('timing') if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c['main'].as_int('row_limit') self.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') c_dest_warning = c['main'].as_bool('destructive_warning') self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c['main'].as_bool('less_chatty') self.null_string = c['main'].get('null_string', '<null>') self.prompt_format = prompt if prompt is not None else c['main'].get('prompt', self.default_prompt) self.prompt_dsn_format = prompt_dsn self.on_error = c['main']['on_error'].upper() self.decimal_format = c['data_formats']['decimal'] self.float_format = c['data_formats']['float'] self.pgspecial.pset_pager(self.config['main'].as_bool( 'enable_pager') and "on" or "off") self.style_output = style_factory_output( self.syntax_style, c['colors']) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') 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': single_connection, 'less_chatty': less_chatty, 'keyword_casing': keyword_casing, } completer = PGCompleter(smart_completion, pgspecial=self.pgspecial, settings=self.settings) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.eventloop = create_eventloop() self.cli = None def quit(self): raise PgCliQuitError def register_special_commands(self): self.pgspecial.register( self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) refresh_callback = lambda: self.refresh_completions( persist_priorities='all') self.pgspecial.register(self.quit, '\\q', '\\q', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=True, aliases=(':q',)) self.pgspecial.register(self.quit, 'quit', 'quit', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=False, aliases=('exit',)) self.pgspecial.register(refresh_callback, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(refresh_callback, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') self.pgspecial.register(self.write_to_file, '\\o', '\\o [filename]', 'Send all query results to file.') self.pgspecial.register(self.info_connection, '\\conninfo', '\\conninfo', 'Get connection details') self.pgspecial.register(self.change_table_format, '\\T', '\\T [format]', 'Change the table format used to output results') def change_table_format(self, pattern, **_): try: if pattern not in TabularOutputFormatter().supported_formats: raise ValueError() self.table_format = pattern yield (None, None, None, 'Changed table format to {}'.format(pattern)) except ValueError: msg = 'Table format {} not recognized. Allowed formats:'.format( pattern) for table_type in TabularOutputFormatter().supported_formats: msg += "\n\t{}".format(table_type) msg += '\nCurrently set to: %s' % self.table_format yield (None, None, None, msg) def info_connection(self, **_): if self.pgexecute.host.startswith('/'): host = 'socket "%s"' % self.pgexecute.host else: host = 'host "%s"' % self.pgexecute.host yield (None, None, None, 'You are connected to database "%s" as user ' '"%s" on %s at port "%s".' % (self.pgexecute.dbname, self.pgexecute.user, host, self.pgexecute.port)) def change_db(self, pattern, **_): if pattern: # Get all the parameters in pattern, handling double quotes if any. infos = re.findall(r'"[^"]*"|[^"\'\s]+', pattern) # Now removing quotes. list(map(lambda s: s.strip('"'), infos)) infos.extend([None] * (4 - len(infos))) db, user, host, port = infos try: self.pgexecute.connect(database=db, user=user, host=host, port=port, application_name='pgcli') except OperationalError as e: click.secho(str(e), err=True, fg='red') click.echo("Previous connection kept") else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message, '', False)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e), '', False)] if (self.destructive_warning and confirm_destructive_query(query) is False): message = 'Wise choice. Command execution stopped.' return [(None, None, None, message)] on_error_resume = (self.on_error == 'RESUME') return self.pgexecute.run( query, self.pgspecial, on_error_resume=on_error_resume ) 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() + '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('pgcli') root_logger.addHandler(handler) root_logger.setLevel(log_level) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) pgspecial_logger = logging.getLogger('pgspecial') pgspecial_logger.addHandler(handler) pgspecial_logger.setLevel(log_level) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash def fixup_possible_percent_encoding(s): return unquote(str(s)) if s else s arguments = dict(database=fixup_possible_percent_encoding(database), host=fixup_possible_percent_encoding(uri.hostname), user=fixup_possible_percent_encoding(uri.username), port=fixup_possible_percent_encoding(uri.port), passwd=fixup_possible_percent_encoding(uri.password)) # Deal with extra params e.g. ?sslmode=verify-ca&sslrootcert=/myrootcert if uri.query: arguments = dict( {k: v for k, (v,) in parse_qs(uri.query).items()}, **arguments) # unquote str(each URI part (they may be percent encoded) self.connect(**arguments) def connect(self, database='', host='', user='', port='', passwd='', dsn='', **kwargs): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Find password from store key = '%s@%s' % (user, host) if not passwd: passwd = keyring.get_password('pgcli', key) # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password for %s' % user, hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn, application_name='pgcli', **kwargs) if passwd: keyring.set_password('pgcli', key, passwd) except (OperationalError, InterfaceError) as e: if ('no password supplied' in utf8tounicode(e.args[0]) and auto_passwd_prompt): passwd = click.prompt('Password for %s' % user, hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn, application_name='pgcli', **kwargs) if passwd: keyring.set_password('pgcli', key, passwd) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute 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 try: 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() finally: cli.application.pre_run_callables = saved_callables return document def execute_command(self, text, query): logger = self.logger try: if (self.destructive_warning): destroy = confirm = confirm_destructive_query(text) if destroy is None: output, query = self._evaluate_command(text) elif destroy is True: click.secho('Your call!') output, query = self._evaluate_command(text) else: click.secho('Wise choice!') raise KeyboardInterrupt except KeyboardInterrupt: # Restart connection to the database self.pgexecute.connect() logger.debug("cancelled query, sql: %r", text) click.secho("cancelled query", err=True, fg='red') except NotImplementedError: click.secho('Not Yet Implemented.', fg="yellow") except OperationalError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self._handle_server_closed_connection() except PgCliQuitError as e: raise 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: self.echo_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second 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') elif query.path_changed: logger.debug('Refreshing search path') with self._completer_lock: self.completer.set_search_path( self.pgexecute.search_path()) logger.debug('Search path: %r', self.completer.search_path) return query def run_cli(self): logger = self.logger history_file = self.config['main']['history_file'] if history_file == 'default': history_file = config_location() + 'history' history = FileHistory(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:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') try: while True: document = self.cli.run() 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()) click.secho(str(e), err=True, fg='red') continue # Initialize default metaquery in case execution fails query = MetaQuery(query=document.text, successful=False) self.watch_command, timing = special.get_watch_command( document.text) if self.watch_command: while self.watch_command: try: query = self.execute_command( self.watch_command, query) click.echo( 'Waiting for {0} seconds before repeating' .format(timing)) sleep(timing) except KeyboardInterrupt: self.watch_command = None else: query = self.execute_command(document.text, query) self.now = dt.datetime.today() # Allow PGCompleter to learn user's preferred keywords, etc. with self._completer_lock: self.completer.extend_query_history(document.text) self.query_history.append(query) except PgCliQuitError: if not self.less_chatty: print ('Goodbye!') 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 _should_show_limit_prompt(self, status, cur): """returns True if limit prompt should be shown, False otherwise.""" if not is_select(status): return False return self.row_limit > 0 and cur and cur.rowcount > self.row_limit def _evaluate_command(self, text): """Used to run a command entered by the user during CLI operation (Puts the E in REPL) returns (results, MetaQuery) """ logger = self.logger logger.debug('sql: %r', text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 # Run the query. start = time() on_error_resume = self.on_error == 'RESUME' res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) threshold = self.row_limit if self._should_show_limit_prompt(status, cur): click.secho('The result set has more than %s rows.' % threshold, fg='red') if not click.confirm('Do you want to continue?'): click.secho("Aborted!", err=True, fg='red') break if self.pgspecial.auto_expand or self.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None expanded = self.pgspecial.expanded_output or self.expanded_output settings = OutputSettings( table_format=self.table_format, dcmlfmt=self.decimal_format, floatfmt=self.float_format, missingval=self.null_string, expanded=expanded, max_width=max_width, case_function=( self.completer.case if self.settings['case_column_headers'] else lambda x: x ), style_output=self.style_output ) formatted = format_output(title, cur, headers, status, settings) output.extend(formatted) total = time() - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, meta_changed, db_changed, path_changed, mutated) return output, meta_query 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.pgexecute.connect() click.secho('Reconnected!\nTry the command again.', fg='green') except OperationalError as e: click.secho(str(e), err=True, fg='red') def refresh_completions(self, history=None, persist_priorities='all'): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial(self._on_completions_refreshed, persist_priorities=persist_priorities) self.completion_refresher.refresh(self.pgexecute, self.pgspecial, 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 pgcli 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): # should be before replacing \\d string = string.replace('\\dsn_alias', self.dsn_alias or '') string = string.replace('\\t', self.now.strftime('%x %X')) string = string.replace('\\u', self.pgexecute.user or '(none)') host = self.pgexecute.host or '(none)' string = string.replace('\\H', host) short_host, _, _ = host.partition('.') string = string.replace('\\h', short_host) string = string.replace('\\d', self.pgexecute.dbname or '(none)') string = string.replace('\\p', str(self.pgexecute.port) or '(none)') string = string.replace('\\i', str(self.pgexecute.pid) or '(none)') string = string.replace('\\#', "#" if (self.pgexecute.superuser) else ">") 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 echo_via_pager(self, text, color=None): if self.pgspecial.pager_config == PAGER_OFF or self.watch_command: click.echo(text, color=color) else: click.echo_via_pager(text, color)
class PGCli(object): def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) default_config = os.path.join(package_root, 'pgclirc') write_default_config(default_config, pgclirc_file) self.pgspecial = PGSpecial() # Load config. c = self.config = load_config(pgclirc_file, default_config) self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') 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.completion_refresher = CompletionRefresher() self.logger = logging.getLogger(__name__) self.initialize_logging() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.cli = None def register_special_commands(self): self.pgspecial.register(self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) self.pgspecial.register(self.refresh_completions, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.refresh_completions, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') def change_db(self, pattern, **_): if pattern: db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern self.pgexecute.connect(database=db) else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e))] return self.pgexecute.run(query, self.pgspecial) def initialize_logging(self): log_file = self.config['main']['log_file'] log_level = self.config['main']['log_level'] level_map = { 'CRITICAL': logging.CRITICAL, 'ERROR': logging.ERROR, 'WARNING': logging.WARNING, 'INFO': logging.INFO, 'DEBUG': logging.DEBUG } handler = logging.FileHandler(os.path.expanduser(log_file)) formatter = logging.Formatter( '%(asctime)s (%(process)d/%(threadName)s) ' '%(name)s %(levelname)s - %(message)s') handler.setFormatter(formatter) root_logger = logging.getLogger('pgcli') root_logger.addHandler(handler) root_logger.setLevel(level_map[log_level.upper()]) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash self.connect(database, uri.hostname, uri.username, uri.port, uri.password) def connect(self, database='', host='', user='', port='', passwd='', dsn=''): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn) except OperationalError as e: if ('no password supplied' in utf8tounicode(e.args[0]) and auto_passwd_prompt): passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute def handle_editor_command(self, cli, document): """ 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 """ while special.editor_command(document.text): filename = special.get_filename(document.text) sql, message = special.open_external_editor(filename, sql=document.text) if message: # Something went wrong. Raise an exception and bail. raise RuntimeError(message) cli.current_buffer.document = Document(sql, cursor_position=len(sql)) document = cli.run(False) continue return document def run_cli(self): pgexecute = self.pgexecute logger = self.logger original_less_opts = self.adjust_less_opts() 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, lambda: self.completion_refresher.is_refreshing()) layout = create_default_layout( lexer=PostgresLexer, reserve_space_for_menu=True, get_prompt_tokens=prompt_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()), ]) history_file = self.config['main']['history_file'] with self._completer_lock: buf = PGBuffer(always_multiline=self.multi_line, completer=self.completer, history=FileHistory( os.path.expanduser(history_file)), 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) self.cli = CommandLineInterface(application=application, eventloop=create_eventloop()) try: while True: document = self.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(self.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 = [] # Run the query. start = time() res = pgexecute.run(document.text, self.pgspecial) 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): 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 if self.pgspecial.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output( title, cur, headers, status, self.table_format, self.pgspecial.expanded_output, max_width) 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: try: click.echo_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: print('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') with self._completer_lock: self.completer.set_search_path(pgexecute.search_path()) logger.debug('Search path: %r', self.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 adjust_less_opts(self): less_opts = os.environ.get('LESS', '') self.logger.debug('Original value for LESS env var: %r', less_opts) os.environ['LESS'] = '-SRXF' return less_opts def refresh_completions(self): self.completion_refresher.refresh(self.pgexecute, self.pgspecial, self._on_completions_refreshed) return [(None, None, None, 'Auto-completion refresh started in the background.')] def _on_completions_refreshed(self, new_completer): self._swap_completer_objects(new_completer) 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): """Swap the completer object in cli with the newly created completer. """ with self._completer_lock: self.completer = new_completer # When pgcli 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 pgspecial(): return PGSpecial()
def __init__( self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None, ): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None self.watch_command = None # Load config. c = self.config = get_config(pgclirc_file) NamedQueries.instance = NamedQueries.from_config(self.config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c["main"].as_bool("multi_line") self.multiline_mode = c["main"].get("multi_line_mode", "psql") self.vi_mode = c["main"].as_bool("vi") self.auto_expand = auto_vertical_output or c["main"].as_bool("auto_expand") self.expanded_output = c["main"].as_bool("expand") self.pgspecial.timing_enabled = c["main"].as_bool("timing") if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c["main"].as_int("row_limit") self.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") c_dest_warning = c["main"].as_bool("destructive_warning") self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c["main"].as_bool("less_chatty") self.null_string = c["main"].get("null_string", "<null>") self.prompt_format = ( prompt if prompt is not None else c["main"].get("prompt", self.default_prompt) ) self.prompt_dsn_format = prompt_dsn self.on_error = c["main"]["on_error"].upper() self.decimal_format = c["data_formats"]["decimal"] self.float_format = c["data_formats"]["float"] self.initialize_keyring() self.pgspecial.pset_pager( self.config["main"].as_bool("enable_pager") and "on" or "off" ) self.style_output = style_factory_output(self.syntax_style, c["colors"]) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c["main"].as_bool("smart_completion") 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": single_connection, "less_chatty": less_chatty, "keyword_casing": keyword_casing, } completer = PGCompleter( smart_completion, pgspecial=self.pgspecial, settings=self.settings ) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.prompt_app = None
class PGCli(object): default_prompt = "\\u@\\h:\\d> " max_len_prompt = 30 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: "%s"', configured_pager ) os.environ["PAGER"] = configured_pager elif os_environ_pager: self.logger.info( 'Default pager found in PAGER environment variable: "%s"', 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, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None, row_limit=None, single_connection=False, less_chatty=None, prompt=None, prompt_dsn=None, auto_vertical_output=False, warn=None, ): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute self.dsn_alias = None self.watch_command = None # Load config. c = self.config = get_config(pgclirc_file) NamedQueries.instance = NamedQueries.from_config(self.config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.output_file = None self.pgspecial = PGSpecial() self.multi_line = c["main"].as_bool("multi_line") self.multiline_mode = c["main"].get("multi_line_mode", "psql") self.vi_mode = c["main"].as_bool("vi") self.auto_expand = auto_vertical_output or c["main"].as_bool("auto_expand") self.expanded_output = c["main"].as_bool("expand") self.pgspecial.timing_enabled = c["main"].as_bool("timing") if row_limit is not None: self.row_limit = row_limit else: self.row_limit = c["main"].as_int("row_limit") self.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") c_dest_warning = c["main"].as_bool("destructive_warning") self.destructive_warning = c_dest_warning if warn is None else warn self.less_chatty = bool(less_chatty) or c["main"].as_bool("less_chatty") self.null_string = c["main"].get("null_string", "<null>") self.prompt_format = ( prompt if prompt is not None else c["main"].get("prompt", self.default_prompt) ) self.prompt_dsn_format = prompt_dsn self.on_error = c["main"]["on_error"].upper() self.decimal_format = c["data_formats"]["decimal"] self.float_format = c["data_formats"]["float"] self.initialize_keyring() self.pgspecial.pset_pager( self.config["main"].as_bool("enable_pager") and "on" or "off" ) self.style_output = style_factory_output(self.syntax_style, c["colors"]) self.now = dt.datetime.today() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c["main"].as_bool("smart_completion") 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": single_connection, "less_chatty": less_chatty, "keyword_casing": keyword_casing, } completer = PGCompleter( smart_completion, pgspecial=self.pgspecial, settings=self.settings ) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.prompt_app = None def quit(self): raise PgCliQuitError def register_special_commands(self): self.pgspecial.register( self.change_db, "\\c", "\\c[onnect] database_name", "Change to a new database.", aliases=("use", "\\connect", "USE"), ) refresh_callback = lambda: self.refresh_completions(persist_priorities="all") self.pgspecial.register( self.quit, "\\q", "\\q", "Quit pgcli.", arg_type=NO_QUERY, case_sensitive=True, aliases=(":q",), ) self.pgspecial.register( self.quit, "quit", "quit", "Quit pgcli.", arg_type=NO_QUERY, case_sensitive=False, aliases=("exit",), ) self.pgspecial.register( refresh_callback, "\\#", "\\#", "Refresh auto-completions.", arg_type=NO_QUERY, ) self.pgspecial.register( refresh_callback, "\\refresh", "\\refresh", "Refresh auto-completions.", arg_type=NO_QUERY, ) self.pgspecial.register( self.execute_from_file, "\\i", "\\i filename", "Execute commands from file." ) self.pgspecial.register( self.write_to_file, "\\o", "\\o [filename]", "Send all query results to file.", ) self.pgspecial.register( self.info_connection, "\\conninfo", "\\conninfo", "Get connection details" ) self.pgspecial.register( self.change_table_format, "\\T", "\\T [format]", "Change the table format used to output results", ) def change_table_format(self, pattern, **_): try: if pattern not in TabularOutputFormatter().supported_formats: raise ValueError() self.table_format = pattern yield (None, None, None, "Changed table format to {}".format(pattern)) except ValueError: msg = "Table format {} not recognized. Allowed formats:".format(pattern) for table_type in TabularOutputFormatter().supported_formats: msg += "\n\t{}".format(table_type) msg += "\nCurrently set to: %s" % self.table_format yield (None, None, None, msg) def info_connection(self, **_): if self.pgexecute.host.startswith("/"): host = 'socket "%s"' % self.pgexecute.host else: host = 'host "%s"' % self.pgexecute.host yield ( None, None, None, 'You are connected to database "%s" as user ' '"%s" on %s at port "%s".' % (self.pgexecute.dbname, self.pgexecute.user, host, self.pgexecute.port), ) def change_db(self, pattern, **_): if pattern: # Get all the parameters in pattern, handling double quotes if any. infos = re.findall(r'"[^"]*"|[^"\'\s]+', pattern) # Now removing quotes. list(map(lambda s: s.strip('"'), infos)) infos.extend([None] * (4 - len(infos))) db, user, host, port = infos try: self.pgexecute.connect( database=db, user=user, host=host, port=port, **self.pgexecute.extra_args ) except OperationalError as e: click.secho(str(e), err=True, fg="red") click.echo("Previous connection kept") else: self.pgexecute.connect() yield ( None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user), ) def execute_from_file(self, pattern, **_): if not pattern: message = "\\i: missing required argument" return [(None, None, None, message, "", False, True)] try: with open(os.path.expanduser(pattern), encoding="utf-8") as f: query = f.read() except IOError as e: return [(None, None, None, str(e), "", False, True)] if self.destructive_warning and confirm_destructive_query(query) is False: message = "Wise choice. Command execution stopped." return [(None, None, None, message)] on_error_resume = self.on_error == "RESUME" return self.pgexecute.run( query, self.pgspecial, on_error_resume=on_error_resume ) def write_to_file(self, pattern, **_): if not pattern: self.output_file = None message = "File output disabled" return [(None, None, None, message, "", True, 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, True)] self.output_file = filename message = 'Writing to file "%s"' % self.output_file return [(None, None, None, message, "", True, True)] def initialize_logging(self): log_file = self.config["main"]["log_file"] if log_file == "default": log_file = config_location() + "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("pgcli") root_logger.addHandler(handler) root_logger.setLevel(log_level) root_logger.debug("Initializing pgcli logging.") root_logger.debug("Log file %r.", log_file) pgspecial_logger = logging.getLogger("pgspecial") pgspecial_logger.addHandler(handler) pgspecial_logger.setLevel(log_level) def initialize_keyring(self): global keyring keyring_enabled = self.config["main"].as_bool("keyring") if keyring_enabled: # Try best to load keyring (issue #1041). import importlib try: keyring = importlib.import_module("keyring") except Exception as e: # ImportError for Python 2, ModuleNotFoundError for Python 3 self.logger.warning("import keyring failed: %r.", e) def connect_dsn(self, dsn, **kwargs): self.connect(dsn=dsn, **kwargs) def connect_uri(self, uri): kwargs = psycopg2.extensions.parse_dsn(uri) remap = {"dbname": "database", "password": "******"} kwargs = {remap.get(k, k): v for k, v in kwargs.items()} self.connect(**kwargs) def connect( self, database="", host="", user="", port="", passwd="", dsn="", **kwargs ): # Connect to the database. if not user: user = getuser() if not database: database = user kwargs.setdefault("application_name", "pgcli") # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get("PGPASSWORD", "") # Find password from store key = "%s@%s" % (user, host) keyring_error_message = dedent( """\ {} {} To remove this message do one of the following: - prepare keyring as described at: https://keyring.readthedocs.io/en/stable/ - uninstall keyring: pip uninstall keyring - disable keyring in our configuration: add keyring = False to [main]""" ) if not passwd and keyring: try: passwd = keyring.get_password("pgcli", key) except (RuntimeError, keyring.errors.InitError) as e: click.secho( keyring_error_message.format( "Load your password from keyring returned:", str(e) ), err=True, fg="red", ) # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt( "Password for %s" % user, hide_input=True, show_default=False, type=str ) def should_ask_for_password(exc): # Prompt for a password after 1st attempt to connect # fails. Don't prompt if the -w flag is supplied if self.never_passwd_prompt: return False error_msg = utf8tounicode(exc.args[0]) if "no password supplied" in error_msg: return True if "password authentication failed" in error_msg: return True return False # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing or incorrect password, but we're allowed to # prompt for a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn, **kwargs) except (OperationalError, InterfaceError) as e: if should_ask_for_password(e): passwd = click.prompt( "Password for %s" % user, hide_input=True, show_default=False, type=str, ) pgexecute = PGExecute( database, user, passwd, host, port, dsn, **kwargs ) else: raise e if passwd and keyring: try: keyring.set_password("pgcli", key, passwd) except (RuntimeError, keyring.errors.KeyringError) as e: click.secho( keyring_error_message.format( "Set password in keyring returned:", str(e) ), err=True, fg="red", ) except Exception as e: # Connecting to a database could fail. self.logger.debug("Database connection failed: %r.", e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg="red") exit(1) self.pgexecute = pgexecute 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 """ editor_command = special.editor_command(text) while editor_command: if editor_command == "\\e": filename = special.get_filename(text) query = special.get_editor_query(text) or self.get_last_query() else: # \ev or \ef filename = None spec = text.split()[1] if editor_command == "\\ev": query = self.pgexecute.view_definition(spec) elif editor_command == "\\ef": query = self.pgexecute.function_definition(spec) 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_app.prompt(default=sql) break except KeyboardInterrupt: sql = "" editor_command = special.editor_command(text) return text def execute_command(self, text): logger = self.logger query = MetaQuery(query=text, successful=False) try: if self.destructive_warning: destroy = confirm = confirm_destructive_query(text) if destroy is False: click.secho("Wise choice!") raise KeyboardInterrupt elif destroy: click.secho("Your call!") output, query = self._evaluate_command(text) except KeyboardInterrupt: # Restart connection to the database self.pgexecute.connect() logger.debug("cancelled query, sql: %r", text) click.secho("cancelled query", err=True, fg="red") except NotImplementedError: click.secho("Not Yet Implemented.", fg="yellow") except OperationalError as e: logger.error("sql: %r, error: %r", text, e) logger.error("traceback: %r", traceback.format_exc()) self._handle_server_closed_connection(text) except (PgCliQuitError, EOFError) as e: raise 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: self.echo_via_pager("\n".join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second if query.total_time > 1: print( "Time: %0.03fs (%s), executed in: %0.03fs (%s)" % ( query.total_time, humanize.time.naturaldelta(query.total_time), query.execution_time, humanize.time.naturaldelta(query.execution_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") elif query.path_changed: logger.debug("Refreshing search path") with self._completer_lock: self.completer.set_search_path(self.pgexecute.search_path()) logger.debug("Search path: %r", self.completer.search_path) return query def run_cli(self): logger = self.logger history_file = self.config["main"]["history_file"] if history_file == "default": history_file = config_location() + "history" history = FileHistory(os.path.expanduser(history_file)) self.refresh_completions(history=history, persist_priorities="none") self.prompt_app = self._build_cli(history) if not self.less_chatty: print("Server: PostgreSQL", self.pgexecute.server_version) print("Version:", __version__) print("Chat: https://gitter.im/dbcli/pgcli") print("Home: http://pgcli.com") try: while True: try: text = self.prompt_app.prompt() except KeyboardInterrupt: continue 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()) click.secho(str(e), err=True, fg="red") continue # Initialize default metaquery in case execution fails self.watch_command, timing = special.get_watch_command(text) if self.watch_command: while self.watch_command: try: query = self.execute_command(self.watch_command) click.echo( "Waiting for {0} seconds before repeating".format( timing ) ) sleep(timing) except KeyboardInterrupt: self.watch_command = None else: query = self.execute_command(text) self.now = dt.datetime.today() # Allow PGCompleter to learn user's preferred keywords, etc. with self._completer_lock: self.completer.extend_query_history(text) self.query_history.append(query) except (PgCliQuitError, EOFError): if not self.less_chatty: print("Goodbye!") 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 _should_limit_output(self, sql, cur): """returns True if the output should be truncated, False otherwise.""" if not is_select(sql): return False return ( not self._has_limit(sql) and self.row_limit != 0 and cur and cur.rowcount > self.row_limit ) def _has_limit(self, sql): if not sql: return False return "limit " in sql.lower() def _limit_output(self, cur): limit = min(self.row_limit, cur.rowcount) new_cur = itertools.islice(cur, limit) new_status = "SELECT " + str(limit) click.secho("The result was limited to %s rows" % limit, fg="red") return new_cur, new_status 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) """ logger = self.logger logger.debug("sql: %r", text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 execution = 0 # Run the query. start = time() on_error_resume = self.on_error == "RESUME" res = self.pgexecute.run( text, self.pgspecial, exception_formatter, on_error_resume ) is_special = None for title, cur, headers, status, sql, success, is_special in res: logger.debug("headers: %r", headers) logger.debug("rows: %r", cur) logger.debug("status: %r", status) if self._should_limit_output(sql, cur): cur, status = self._limit_output(cur) if self.pgspecial.auto_expand or self.auto_expand: max_width = self.prompt_app.output.get_size().columns else: max_width = None expanded = self.pgspecial.expanded_output or self.expanded_output settings = OutputSettings( table_format=self.table_format, dcmlfmt=self.decimal_format, floatfmt=self.float_format, missingval=self.null_string, expanded=expanded, max_width=max_width, case_function=( self.completer.case if self.settings["case_column_headers"] else lambda x: x ), style_output=self.style_output, ) execution = time() - start formatted = format_output(title, cur, headers, status, settings) output.extend(formatted) total = time() - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery( text, all_success, total, execution, meta_changed, db_changed, path_changed, mutated, is_special, ) return output, meta_query def _handle_server_closed_connection(self, text): """Used during CLI execution.""" try: click.secho("Reconnecting...", fg="green") self.pgexecute.connect() click.secho("Reconnected!", fg="green") self.execute_command(text) except OperationalError as e: click.secho("Reconnect Failed", fg="red") click.secho(str(e), err=True, fg="red") def refresh_completions(self, history=None, persist_priorities="all"): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial( self._on_completions_refreshed, persist_priorities=persist_priorities ) self.completion_refresher.refresh( self.pgexecute, self.pgspecial, 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_app: # After refreshing, redraw the CLI to clear the statusbar # "Refreshing completions..." indicator self.prompt_app.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 self.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): # should be before replacing \\d string = string.replace("\\dsn_alias", self.dsn_alias or "") string = string.replace("\\t", self.now.strftime("%x %X")) string = string.replace("\\u", self.pgexecute.user or "(none)") string = string.replace("\\H", self.pgexecute.host or "(none)") string = string.replace("\\h", self.pgexecute.short_host or "(none)") string = string.replace("\\d", self.pgexecute.dbname or "(none)") string = string.replace( "\\p", str(self.pgexecute.port) if self.pgexecute.port is not None else "5432", ) string = string.replace("\\i", str(self.pgexecute.pid) or "(none)") string = string.replace("\\#", "#" if (self.pgexecute.superuser) else ">") 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 is_too_wide(self, line): """Will this line be too wide to fit into terminal?""" if not self.prompt_app: return False return ( len(COLOR_CODE_REGEX.sub("", line)) > self.prompt_app.output.get_size().columns ) def is_too_tall(self, lines): """Are there too many lines to fit into terminal?""" if not self.prompt_app: return False return len(lines) >= (self.prompt_app.output.get_size().rows - 4) def echo_via_pager(self, text, color=None): if self.pgspecial.pager_config == PAGER_OFF or self.watch_command: click.echo(text, color=color) elif self.pgspecial.pager_config == PAGER_LONG_OUTPUT: lines = text.split("\n") # The last 4 lines are reserved for the pgcli menu and padding if self.is_too_tall(lines) or any(self.is_too_wide(l) for l in lines): click.echo_via_pager(text, color=color) else: click.echo(text, color=color) else: click.echo_via_pager(text, color)
class PGCli(object): 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: ' + '\'' + configured_pager + '\'') os.environ['PAGER'] = configured_pager elif (os_environ_pager): self.logger.info( 'Default pager found in PAGER environment variable: ' + '\'' + os_environ_pager + '\'') os.environ['PAGER'] = os_environ_pager else: self.logger.info( 'No default pager found in environment. Using os default pager' ) # Always set default set of less recommended options, they are ignored if pager is # different than less or is already parameterized with their own arguments os.environ['LESS'] = '-SRXF' def __init__(self, force_passwd_prompt=False, never_passwd_prompt=False, pgexecute=None, pgclirc_file=None): self.force_passwd_prompt = force_passwd_prompt self.never_passwd_prompt = never_passwd_prompt self.pgexecute = pgexecute from pgcli import __file__ as package_root package_root = os.path.dirname(package_root) default_config = os.path.join(package_root, 'pgclirc') write_default_config(default_config, pgclirc_file) # Load config. c = self.config = load_config(pgclirc_file, default_config) self.logger = logging.getLogger(__name__) self.initialize_logging() self.set_default_pager(c) self.pgspecial = PGSpecial() self.multi_line = c['main'].as_bool('multi_line') self.vi_mode = c['main'].as_bool('vi') self.pgspecial.timing_enabled = c['main'].as_bool('timing') 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.on_error = c['main']['on_error'].upper() self.completion_refresher = CompletionRefresher() self.query_history = [] # Initialize completer smart_completion = c['main'].as_bool('smart_completion') completer = PGCompleter(smart_completion, pgspecial=self.pgspecial) self.completer = completer self._completer_lock = threading.Lock() self.register_special_commands() self.cli = None def register_special_commands(self): self.pgspecial.register(self.change_db, '\\c', '\\c[onnect] database_name', 'Change to a new database.', aliases=('use', '\\connect', 'USE')) refresh_callback = lambda: self.refresh_completions(persist_priorities= 'all') self.pgspecial.register(refresh_callback, '\\#', '\\#', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(refresh_callback, '\\refresh', '\\refresh', 'Refresh auto-completions.', arg_type=NO_QUERY) self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename', 'Execute commands from file.') def change_db(self, pattern, **_): if pattern: db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern self.pgexecute.connect(database=db) else: self.pgexecute.connect() yield (None, None, None, 'You are now connected to database "%s" as ' 'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user)) def execute_from_file(self, pattern, **_): if not pattern: message = '\\i: missing required argument' return [(None, None, None, message, '', False)] try: with open(os.path.expanduser(pattern), encoding='utf-8') as f: query = f.read() except IOError as e: return [(None, None, None, str(e), '', False)] on_error_resume = (self.on_error == 'RESUME') return self.pgexecute.run(query, self.pgspecial, on_error_resume=on_error_resume) def initialize_logging(self): log_file = self.config['main']['log_file'] if log_file == 'default': log_file = config_location() + 'log' ensure_dir_exists(log_file) log_level = self.config['main']['log_level'] level_map = { 'CRITICAL': logging.CRITICAL, 'ERROR': logging.ERROR, 'WARNING': logging.WARNING, 'INFO': logging.INFO, 'DEBUG': logging.DEBUG } handler = logging.FileHandler(os.path.expanduser(log_file)) formatter = logging.Formatter( '%(asctime)s (%(process)d/%(threadName)s) ' '%(name)s %(levelname)s - %(message)s') handler.setFormatter(formatter) root_logger = logging.getLogger('pgcli') root_logger.addHandler(handler) root_logger.setLevel(level_map[log_level.upper()]) root_logger.debug('Initializing pgcli logging.') root_logger.debug('Log file %r.', log_file) def connect_dsn(self, dsn): self.connect(dsn=dsn) def connect_uri(self, uri): uri = urlparse(uri) database = uri.path[1:] # ignore the leading fwd slash self.connect(database, uri.hostname, uri.username, uri.port, uri.password) def connect(self, database='', host='', user='', port='', passwd='', dsn=''): # Connect to the database. if not user: user = getuser() if not database: database = user # If password prompt is not forced but no password is provided, try # getting it from environment variable. if not self.force_passwd_prompt and not passwd: passwd = os.environ.get('PGPASSWORD', '') # Prompt for a password immediately if requested via the -W flag. This # avoids wasting time trying to connect to the database and catching a # no-password exception. # If we successfully parsed a password from a URI, there's no need to # prompt for it, even with the -W flag if self.force_passwd_prompt and not passwd: passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) # Prompt for a password after 1st attempt to connect without a password # fails. Don't prompt if the -w flag is supplied auto_passwd_prompt = not passwd and not self.never_passwd_prompt # Attempt to connect to the database. # Note that passwd may be empty on the first attempt. If connection # fails because of a missing password, but we're allowed to prompt for # a password (no -w flag), prompt for a passwd and try again. try: try: pgexecute = PGExecute(database, user, passwd, host, port, dsn) except OperationalError as e: if ('no password supplied' in utf8tounicode(e.args[0]) and auto_passwd_prompt): passwd = click.prompt('Password', hide_input=True, show_default=False, type=str) pgexecute = PGExecute(database, user, passwd, host, port, dsn) else: raise e except Exception as e: # Connecting to a database could fail. self.logger.debug('Database connection failed: %r.', e) self.logger.error("traceback: %r", traceback.format_exc()) click.secho(str(e), err=True, fg='red') exit(1) self.pgexecute = pgexecute def handle_editor_command(self, cli, document): """ 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 """ while special.editor_command(document.text): filename = special.get_filename(document.text) sql, message = special.open_external_editor(filename, sql=document.text) if message: # Something went wrong. Raise an exception and bail. raise RuntimeError(message) cli.current_buffer.document = Document(sql, cursor_position=len(sql)) document = cli.run(False) continue return document def run_cli(self): logger = self.logger history_file = self.config['main']['history_file'] if history_file == 'default': history_file = config_location() + 'history' history = FileHistory(os.path.expanduser(history_file)) self.refresh_completions(history=history, persist_priorities='none') self.cli = self._build_cli(history) print('Version:', __version__) print('Chat: https://gitter.im/dbcli/pgcli') print('Mail: https://groups.google.com/forum/#!forum/pgcli') print('Home: http://pgcli.com') try: while True: document = self.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(self.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 # Initialize default metaquery in case execution fails query = MetaQuery(query=document.text, successful=False) try: output, query = self._evaluate_command(document.text) except KeyboardInterrupt: # Restart connection to the database self.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: if ('server closed the connection' in utf8tounicode(e.args[0])): self._handle_server_closed_connection() 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: try: click.echo_via_pager('\n'.join(output)) except KeyboardInterrupt: pass if self.pgspecial.timing_enabled: # Only add humanized time display if > 1 second 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') elif query.path_changed: logger.debug('Refreshing search path') with self._completer_lock: self.completer.set_search_path( self.pgexecute.search_path()) logger.debug('Search path: %r', self.completer.search_path) # Allow PGCompleter 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: print('Goodbye!') 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(_): return [(Token.Prompt, '%s> ' % self.pgexecute.dbname)] def get_continuation_tokens(cli, width): return [(Token.Continuation, '.' * (width - 1) + ' ')] get_toolbar_tokens = create_toolbar_tokens_func( lambda: self.vi_mode, self.completion_refresher.is_refreshing) layout = create_prompt_layout( lexer=PygmentsLexer(PostgresLexer), reserve_space_for_menu=4, 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(always_multiline=self.multi_line, completer=self.completer, history=history, complete_while_typing=Always(), accept_action=AcceptAction.RETURN_DOCUMENT) 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) cli = CommandLineInterface(application=application) return cli 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) """ logger = self.logger logger.debug('sql: %r', text) all_success = True meta_changed = False # CREATE, ALTER, DROP, etc mutated = False # INSERT, DELETE, etc db_changed = False path_changed = False output = [] total = 0 # Run the query. start = time() on_error_resume = self.on_error == 'RESUME' res = self.pgexecute.run(text, self.pgspecial, exception_formatter, on_error_resume) for title, cur, headers, status, sql, success 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): 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 if self.pgspecial.auto_expand: max_width = self.cli.output.get_size().columns else: max_width = None formatted = format_output(title, cur, headers, status, self.table_format, self.pgspecial.expanded_output, max_width) output.extend(formatted) end = time() total += end - start # Keep track of whether any of the queries are mutating or changing # the database if success: mutated = mutated or is_mutating(status) db_changed = db_changed or has_change_db_cmd(sql) meta_changed = meta_changed or has_meta_cmd(sql) path_changed = path_changed or has_change_path_cmd(sql) else: all_success = False meta_query = MetaQuery(text, all_success, total, meta_changed, db_changed, path_changed, mutated) return output, meta_query 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.pgexecute.connect() click.secho('Reconnected!\nTry the command again.', fg='green') except OperationalError as e: click.secho(str(e), err=True, fg='red') def refresh_completions(self, history=None, persist_priorities='all'): """ Refresh outdated completions :param history: A prompt_toolkit.history.FileHistory object. Used to load keyword and identifier preferences :param persist_priorities: 'all' or 'keywords' """ callback = functools.partial(self._on_completions_refreshed, persist_priorities=persist_priorities) self.completion_refresher.refresh(self.pgexecute, self.pgspecial, callback, history=history) 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 pgcli 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)