def run(self):
eventloop = create_eventloop()
try:
cli = CommandLineInterface(application=self.application,
eventloop=eventloop)
cli.run()
finally:
eventloop.close()
class AqPrompt(object):
def __init__(self, parser, engine, options=None):
self.parser = parser
self.engine = engine
self.options = options if options is not None else {}
util.ensure_data_dir_exists()
application = create_prompt_application(
message='> ',
lexer=PygmentsLexer(SqlLexer),
history=FileHistory(os.path.expanduser('~/.aq/history')),
completer=AqCompleter(schemas=engine.available_schemas, tables=engine.available_tables),
auto_suggest=AutoSuggestFromHistory(),
validator=QueryValidator(parser),
on_abort=AbortAction.RETRY,
)
loop = create_eventloop()
self.cli = CommandLineInterface(application=application, eventloop=loop)
self.patch_context = self.cli.patch_stdout_context()
def prompt(self):
with self.patch_context:
return self.cli.run(reset_current_buffer=True).text
def update_with_result(self, query_metadata):
# TODO
pass
def run():
validate_osenvironment()
print_banner()
cli_buffer = OSBuffer()
ltobj = OSLayout(multiwindow=False)
application = Application(style=PygmentsStyle(OSStyle),
layout=ltobj.layout,
buffers=cli_buffer.buffers,
on_exit=AbortAction.RAISE_EXCEPTION,
key_bindings_registry=OSKeyBinder.registry)
cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
while True:
try:
document = cli.run(reset_current_buffer=True)
process_document(document)
except KeyboardInterrupt:
# A keyboardInterrupt generated possibly due to Ctrl-C
print "Keyboard Interrupt Generated"
continue
except EOFError:
print "cntl-D"
sys.exit()
class AqPrompt(object):
def __init__(self, parser, engine, options=None):
self.parser = parser
self.engine = engine
self.options = options if options is not None else {}
util.ensure_data_dir_exists()
application = create_prompt_application(
message='> ',
lexer=PygmentsLexer(SqlLexer),
history=FileHistory(os.path.expanduser('~/.aq/history')),
completer=AqCompleter(schemas=engine.available_schemas,
tables=engine.available_tables),
auto_suggest=AutoSuggestFromHistory(),
validator=QueryValidator(parser),
on_abort=AbortAction.RETRY,
)
loop = create_eventloop()
self.cli = CommandLineInterface(application=application,
eventloop=loop)
self.patch_context = self.cli.patch_stdout_context()
def prompt(self):
with self.patch_context:
return self.cli.run(reset_current_buffer=True).text
def update_with_result(self, query_metadata):
# TODO
pass
def run(self):
labels = self.neo4j.get_labels()
relationship_types = self.neo4j.get_relationship_types()
properties = self.neo4j.get_property_keys()
if self.filename:
with open(self.filename, "rb") as f:
queries = split_queries_on_semicolons(f.read())
for query in queries:
print("> " + query)
self.handle_query(query)
print()
return
click.secho(" ______ __ __ ______ __ __ ", fg="red")
click.secho("/\ ___\ /\ \_\ \ /\ ___\ /\ \ /\ \ ", fg="yellow")
click.secho("\ \ \____ \ \____ \ \ \ \____ \ \ \____ \ \ \ ", fg="green")
click.secho(" \ \_____\ \/\_____\ \ \_____\ \ \_____\ \ \_\ ", fg="blue")
click.secho(" \/_____/ \/_____/ \/_____/ \/_____/ \/_/ ", fg="magenta")
print("Cycli version: {}".format(__version__))
print("Neo4j version: {}".format(".".join(map(str, self.neo4j.neo4j_version))))
print("Bug reports: https://github.com/nicolewhite/cycli/issues\n")
completer = CypherCompleter(labels, relationship_types, properties)
layout = create_prompt_layout(
lexer=CypherLexer,
get_prompt_tokens=get_tokens,
reserve_space_for_menu=8,
)
buff = CypherBuffer(
accept_action=AcceptAction.RETURN_DOCUMENT,
history=FileHistory(filename=os.path.expanduser('~/.cycli_history')),
completer=completer,
complete_while_typing=True,
)
application = Application(
style=PygmentsStyle(CypherStyle),
buffer=buff,
layout=layout,
on_exit=AbortAction.RAISE_EXCEPTION,
key_bindings_registry=CypherBinder.registry
)
cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = cli.run()
query = document.text
self.handle_query(query)
except UserWantsOut:
print("Goodbye!")
except Exception as e:
print(e)
def run(self):
neo4j = Neo4j(self.host, self.port, self.username, self.password)
try:
labels = neo4j.labels()
relationship_types = neo4j.relationship_types()
properties = neo4j.properties()
except Unauthorized:
print("Unauthorized. See cycli --help for authorization instructions.")
return
except SocketError:
print("Connection refused. Is Neo4j turned on?")
return
completer = CypherCompleter(labels, relationship_types, properties)
layout = create_default_layout(
lexer=CypherLexer,
get_prompt_tokens=get_tokens,
reserve_space_for_menu=True
)
buff = CypherBuffer(
history=History(),
completer=completer,
complete_while_typing=Always()
)
application = Application(
style=CypherStyle,
buffer=buff,
layout=layout,
on_exit=AbortAction.RAISE_EXCEPTION
)
cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = cli.run()
query = document.text
if query in ["quit", "exit"]:
raise Exception
elif query == "help":
print(help_text())
else:
results = neo4j.cypher(query)
print(results)
except Exception:
print("Goodbye!")
class SaltCli(object):
"""
The CLI implementation.
"""
def __init__(self):
self.logger = logging.getLogger(__name__)
self.id_completer = WordCompleter([
'id1', 'id2', 'id3'
])
def run_cli(self):
"""
Run the main loop
"""
print(u'Version:', __version__)
print(u'Home: https://github.com/glasslion/saltcli')
history = FileHistory(os.path.expanduser('~/.saltcli-history'))
layout = create_prompt_layout(
message=u'saltcli> ',
)
application = Application(
layout=layout
)
eventloop = create_eventloop()
self.cli = CommandLineInterface(
application=application,
eventloop=eventloop)
while True:
try:
document = self.cli.run()
if quit_command(document.text):
raise EOFError
except KeyboardInterrupt:
# user pressed Ctrl + C
click.echo('')
except EOFError:
break
except Exception as ex:
self.logger.debug('Exception: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho("{0}".format(ex), fg='red')
break
print('Goodbye!')
def loop(cmd, history_file):
key_binding_manager = KeyBindingManager(
enable_search=True,
enable_abort_and_exit_bindings=True
)
layout = create_prompt_layout(
message=u'cr> ',
multiline=True,
lexer=SqlLexer,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
buffer = CrashBuffer(
history=TruncatedFileHistory(history_file, max_length=MAX_HISTORY_LENGTH),
accept_action=AcceptAction.RETURN_DOCUMENT,
completer=SQLCompleter(cmd)
)
buffer.complete_while_typing = lambda cli=None: cmd.should_autocomplete()
application = Application(
layout=layout,
buffer=buffer,
style=PygmentsStyle.from_defaults(pygments_style_cls=CrateStyle),
key_bindings_registry=key_binding_manager.registry,
editing_mode=_get_editing_mode(),
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
)
eventloop = create_eventloop()
output = create_output()
cli = CommandLineInterface(
application=application,
eventloop=eventloop,
output=output
)
def get_num_columns_override():
return output.get_size().columns
cmd.get_num_columns = get_num_columns_override
while True:
try:
doc = cli.run(reset_current_buffer=True)
if doc:
cmd.process(doc.text)
except KeyboardInterrupt:
cmd.logger.warn("Query not cancelled. Run KILL <jobId> to cancel it")
except EOFError:
cmd.logger.warn(u'Bye!')
return
class Prompt(object): #pragma: no cover
def __init__(self, renv):
self.renv = renv
self.is_long = True
self.cli = None
def initialize(self):
history = InMemoryHistory()
toolbar_handler = create_toolbar_handler(self.get_long_options)
layout = create_prompt_layout(
get_prompt_tokens=self.get_prompt_tokens,
lexer=create_lexer(),
get_bottom_toolbar_tokens=toolbar_handler)
buf = Buffer(history=history,
completer=CrutchCompleter(self.renv),
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
manager = get_key_manager(self.set_long_options, self.get_long_options)
application = Application(style=style_factory(),
layout=layout,
buffer=buf,
key_bindings_registry=manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.cli = CommandLineInterface(application=application,
eventloop=eventloop)
def get_prompt_tokens(self, _):
return [(Token.Pound, ' Y '), (Token.Text, ' ')]
def set_long_options(self, value):
self.is_long = value
def get_long_options(self):
return self.is_long
def activate(self, reinitialize=False):
if reinitialize:
self.initialize()
assert self.cli
document = self.cli.run(True)
return shlex.split(document.text)
def loop(cmd, history_file):
key_binding_manager = KeyBindingManager(
enable_search=True, enable_abort_and_exit_bindings=True)
layout = create_prompt_layout(
message=u'cr> ',
multiline=True,
lexer=SqlLexer,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
])
buffer = CrashBuffer(history=TruncatedFileHistory(
history_file, max_length=MAX_HISTORY_LENGTH),
accept_action=AcceptAction.RETURN_DOCUMENT,
completer=SQLCompleter(cmd))
buffer.complete_while_typing = lambda cli=None: cmd.should_autocomplete()
application = Application(
layout=layout,
buffer=buffer,
style=PygmentsStyle.from_defaults(pygments_style_cls=CrateStyle),
key_bindings_registry=key_binding_manager.registry,
editing_mode=_get_editing_mode(),
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
)
eventloop = create_eventloop()
output = create_output()
cli = CommandLineInterface(application=application,
eventloop=eventloop,
output=output)
def get_num_columns_override():
return output.get_size().columns
cmd.get_num_columns = get_num_columns_override
while True:
try:
doc = cli.run(reset_current_buffer=True)
if doc:
cmd.process(doc.text)
except KeyboardInterrupt:
cmd.logger.warn(
"Query not cancelled. Run KILL <jobId> to cancel it")
except EOFError:
cmd.logger.warn(u'Bye!')
return
def run():
cli_buffer = VaultBuffer()
vault_layout = VaultLayout(multiwindow=True)
application = Application(style=PygmentsStyle(VaultStyle),
layout=vault_layout.layout,
buffers=cli_buffer.buffers,
on_exit=AbortAction.RAISE_EXCEPTION,
key_bindings_registry=VaultKeyBinder.registry,
use_alternate_screen=False)
cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
while True:
try:
document = cli.run(reset_current_buffer=True)
process_document(document)
except KeyboardInterrupt:
print "Keyboard interrupt generated"
continue
except EOFError:
print "ctrl-D"
sys.exit()
class Saws(object):
"""Encapsulates the Saws CLI.
Attributes:
* aws_cli: An instance of prompt_toolkit's CommandLineInterface.
* key_manager: An instance of KeyManager.
* config: An instance of Config.
* config_obj: An instance of ConfigObj, reads from ~/.sawsrc.
* theme: A string representing the lexer theme.
* logger: An instance of SawsLogger.
* all_commands: A list of all commands, sub_commands, options, etc
from data/SOURCES.txt.
* commands: A list of commands from data/SOURCES.txt.
* sub_commands: A list of sub_commands from data/SOURCES.txt.
* completer: An instance of AwsCompleter.
"""
PYGMENTS_CMD = ' | pygmentize -l json'
def __init__(self, refresh_resources=True):
"""Inits Saws.
Args:
* refresh_resources: A boolean that determines whether to
refresh resources.
Returns:
None.
"""
self.aws_cli = None
self.key_manager = None
self.config = Config()
self.config_obj = self.config.read_configuration()
self.theme = self.config_obj[self.config.MAIN][self.config.THEME]
self.logger = SawsLogger(
__name__,
self.config_obj[self.config.MAIN][self.config.LOG_FILE],
self.config_obj[self.config.MAIN][self.config.LOG_LEVEL]).logger
self.all_commands = AwsCommands().all_commands
self.commands = \
self.all_commands[AwsCommands.CommandType.COMMANDS.value]
self.sub_commands = \
self.all_commands[AwsCommands.CommandType.SUB_COMMANDS.value]
self.completer = AwsCompleter(
awscli_completer,
self.all_commands,
self.config,
self.config_obj,
self.log_exception,
fuzzy_match=self.get_fuzzy_match(),
shortcut_match=self.get_shortcut_match())
if refresh_resources:
self.completer.refresh_resources_and_options()
self._create_cli()
def log_exception(self, e, traceback, echo=False):
"""Logs the exception and traceback to the log file ~/.saws.log.
Args:
* e: A Exception that specifies the exception.
* traceback: A Traceback that specifies the traceback.
* echo: A boolean that specifies whether to echo the exception
to the console using click.
Returns:
None.
"""
self.logger.debug('exception: %r.', str(e))
self.logger.error("traceback: %r", traceback.format_exc())
if echo:
click.secho(str(e), fg='red')
def set_color(self, color):
"""Setter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the color flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.COLOR] = color
def get_color(self):
"""Getter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the color flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.COLOR)
def set_fuzzy_match(self, fuzzy):
"""Setter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the fuzzy flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.FUZZY] = fuzzy
self.completer.fuzzy_match = fuzzy
def get_fuzzy_match(self):
"""Getter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the fuzzy flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.FUZZY)
def set_shortcut_match(self, shortcut):
"""Setter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the shortcut flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.SHORTCUT] = shortcut
self.completer.shortcut_match = shortcut
def get_shortcut_match(self):
"""Getter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the shortcut flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.SHORTCUT)
def refresh_resources_and_options(self):
"""Convenience function to refresh resources and options for completion.
Used by prompt_toolkit's KeyBindingManager.
Args:
* None.
Returns:
None.
"""
self.completer.refresh_resources_and_options(force_refresh=True)
def handle_docs(self, text=None, from_fkey=False):
"""Displays contextual web docs for `F9` or the `docs` command.
Displays the web docs specific to the currently entered:
* (optional) command
* (optional) subcommand
If no command or subcommand is present, the docs index page is shown.
Docs are only displayed if:
* from_fkey is True
* from_fkey is False and `docs` is found in text
Args:
* text: A string representing the input command text.
* from_fkey: A boolean representing whether this function is
being executed from an `F9` key press.
Returns:
A boolean representing whether the web docs were shown.
"""
base_url = 'http://docs.aws.amazon.com/cli/latest/reference/'
index_html = 'index.html'
if text is None:
text = self.aws_cli.current_buffer.document.text
# If the user hit the F9 key, append 'docs' to the text
if from_fkey:
text = text.strip() + ' ' + AwsCommands.AWS_DOCS
tokens = text.split()
if len(tokens) > 2 and tokens[-1] == AwsCommands.AWS_DOCS:
prev_word = tokens[-2]
# If we have a command, build the url
if prev_word in self.commands:
prev_word = prev_word + '/'
url = base_url + prev_word + index_html
webbrowser.open(url)
return True
# if we have a command and subcommand, build the url
elif prev_word in self.sub_commands:
command_url = tokens[-3] + '/'
sub_command_url = tokens[-2] + '.html'
url = base_url + command_url + sub_command_url
webbrowser.open(url)
return True
webbrowser.open(base_url + index_html)
# If we still haven't opened the help doc at this point and the
# user hit the F9 key or typed docs, just open the main docs index
if from_fkey or AwsCommands.AWS_DOCS in tokens:
webbrowser.open(base_url + index_html)
return True
return False
def _handle_cd(self, text):
"""Handles a `cd` shell command by calling python's os.chdir.
Simply passing in the `cd` command to subprocess.call doesn't work.
Note: Changing the directory within Saws will only be in effect while
running Saws. Exiting the program will return you to the directory
you were in prior to running Saws.
Attributes:
* text: A string representing the input command text.
Returns:
A boolean representing a `cd` command was found and handled.
"""
CD_CMD = 'cd'
stripped_text = text.strip()
if stripped_text.startswith(CD_CMD):
directory = ''
if stripped_text == CD_CMD:
# Treat `cd` as a change to the root directory.
# os.path.expanduser does this in a cross platform manner.
directory = os.path.expanduser('~')
else:
tokens = text.split(CD_CMD + ' ')
directory = tokens[-1]
try:
os.chdir(directory)
except OSError as e:
self.log_exception(e, traceback, echo=True)
return True
return False
def _colorize_output(self, text):
"""Highlights output with pygments.
Only highlights the output if all of the following conditions are True:
* The color option is enabled
* The text does not contain the `configure` command
* The text does not contain the `help` command, which already does
output highlighting
Args:
* text: A string that represents the input command text.
Returns:
A string that represents:
* The original command text if no highlighting was performed.
* The pygments highlighted command text otherwise.
"""
stripped_text = text.strip()
if not self.get_color() or stripped_text == '':
return text
excludes = [AwsCommands.AWS_CONFIGURE,
AwsCommands.AWS_HELP,
'|']
if not any(substring in stripped_text for substring in excludes):
return text.strip() + self.PYGMENTS_CMD
else:
return text
def _handle_keyboard_interrupt(self, e, platform):
"""Handles keyboard interrupts more gracefully on Mac/Unix/Linux.
Allows Mac/Unix/Linux to continue running on keyboard interrupt,
as the user might interrupt a long-running AWS command with Control-C
while continuing to work with Saws.
On Windows, the "Terminate batch job (Y/N)" confirmation makes it
tricky to handle this gracefully. Thus, we re-raise KeyboardInterrupt.
Args:
* e: A KeyboardInterrupt.
* platform: A string that denotes platform such as
'Windows', 'Darwin', etc.
Returns:
None
Raises:
Exception: A KeyboardInterrupt if running on Windows.
"""
if platform == 'Windows':
raise e
else:
# Clear the renderer and send a carriage return
self.aws_cli.renderer.clear()
self.aws_cli.input_processor.feed_key(KeyPress(Keys.ControlM, ''))
def _process_command(self, text):
"""Processes the input command, called by the cli event loop
Args:
* text: A string that represents the input command text.
Returns:
None.
"""
if AwsCommands.AWS_COMMAND in text:
text = self.completer.replace_shortcut(text)
if self.handle_docs(text):
return
try:
if not self._handle_cd(text):
text = self._colorize_output(text)
# Pass the command onto the shell so aws-cli can execute it
subprocess.call(text, shell=True)
print('')
except KeyboardInterrupt as e:
self._handle_keyboard_interrupt(e, platform.system())
except Exception as e:
self.log_exception(e, traceback, echo=True)
def _create_cli(self):
"""Creates the prompt_toolkit's CommandLineInterface.
Args:
* None.
Returns:
None.
"""
history = FileHistory(os.path.expanduser('~/.saws-history'))
toolbar = Toolbar(self.get_color,
self.get_fuzzy_match,
self.get_shortcut_match)
layout = create_default_layout(
message='saws> ',
reserve_space_for_menu=True,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar.handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
self.key_manager = KeyManager(
self.set_color,
self.get_color,
self.set_fuzzy_match,
self.get_fuzzy_match,
self.set_shortcut_match,
self.get_shortcut_match,
self.refresh_resources_and_options,
self.handle_docs)
style_factory = StyleFactory(self.theme)
application = Application(
mouse_support=False,
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.aws_cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def run_cli(self):
"""Runs the main loop.
Args:
* None.
Returns:
None.
"""
print('Version:', __version__)
print('Theme:', self.theme)
while True:
document = self.aws_cli.run()
self._process_command(document.text)
def run_cli(self):
pgexecute = self.pgexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
completer = self.completer
self.refresh_completions()
def set_vi_mode(value):
self.vi_mode = value
key_binding_manager = pgcli_bindings(
get_vi_mode_enabled=lambda: self.vi_mode,
set_vi_mode_enabled=set_vi_mode)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/pgcli')
print('Mail: https://groups.google.com/forum/#!forum/pgcli')
print('Home: http://pgcli.com')
def prompt_tokens(cli):
return [(Token.Prompt, '%s> ' % pgexecute.dbname)]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.vi_mode)
layout = create_default_layout(
lexer=PostgresLexer,
reserve_space_for_menu=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
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']
buf = PGBuffer(always_multiline=self.multi_line,
completer=completer,
history=FileHistory(os.path.expanduser(history_file)),
complete_while_typing=Always())
application = Application(
style=style_factory(self.syntax_style, self.cli_style),
layout=layout,
buffer=buf,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
ignore_case=True)
cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = cli.run()
# The reason we check here instead of inside the pgexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the pgexecute.run()
# statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
successful = False
# Initialized to [] because res might never get initialized
# if an exception occurs in pgexecute.run(). Which causes
# finally clause to fail.
res = []
start = time()
# Run the query.
res = pgexecute.run(document.text, self.pgspecial)
duration = time() - start
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and cur
and cur.rowcount > threshold):
click.secho(
'The result set has more than %s rows.' %
threshold,
fg='red')
if not click.confirm('Do you want to continue?'):
click.secho("Aborted!", err=True, fg='red')
break
formatted = format_output(
title, cur, headers, status, self.table_format,
self.pgspecial.expanded_output)
output.extend(formatted)
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
pgexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
click.secho("cancelled query", err=True, fg='red')
except NotImplementedError:
click.secho('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
reconnect = True
if ('server closed the connection'
in utf8tounicode(e.args[0])):
reconnect = click.prompt(
'Connection reset. Reconnect (Y/n)',
show_default=False,
type=bool,
default=True)
if reconnect:
try:
pgexecute.connect()
click.secho(
'Reconnected!\nTry the command again.',
fg='green')
except OperationalError as e:
click.secho(str(e), err=True, fg='red')
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
else:
try:
click.echo_via_pager('\n'.join(output))
except KeyboardInterrupt:
pass
if self.pgspecial.timing_enabled:
print('Command Time: %0.03fs' % duration)
print('Format Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions()
# Refresh search_path to set default schema.
if need_search_path_refresh(document.text):
logger.debug('Refreshing search path')
completer.set_search_path(pgexecute.search_path())
logger.debug('Search path: %r', completer.search_path)
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
print('Goodbye!')
finally: # Reset the less opts back to original.
logger.debug('Restoring env var LESS to %r.', original_less_opts)
os.environ['LESS'] = original_less_opts
def run(self):
labels = self.neo4j.get_labels()
relationship_types = self.neo4j.get_relationship_types()
properties = self.neo4j.get_property_keys()
if self.filename:
queries = self.filename.read()
queries = queries.split(";")[:-1]
for query in queries:
query += ";"
query = query.strip()
print("> " + query)
self.handle_query(query)
print()
return
click.secho(" ______ __ __ ______ __ __ ", fg="red")
click.secho("/\ ___\ /\ \_\ \ /\ ___\ /\ \ /\ \ ", fg="yellow")
click.secho("\ \ \____ \ \____ \ \ \ \____ \ \ \____ \ \ \ ", fg="green")
click.secho(" \ \_____\ \/\_____\ \ \_____\ \ \_____\ \ \_\ ", fg="blue")
click.secho(" \/_____/ \/_____/ \/_____/ \/_____/ \/_/ ", fg="magenta")
print("Cycli version: {}".format(__version__))
print("Neo4j version: {}".format(".".join(map(str, self.neo4j.neo4j_version))))
print("Bug reports: https://github.com/nicolewhite/cycli/issues\n")
completer = CypherCompleter(labels, relationship_types, properties)
layout = create_prompt_layout(
lexer=CypherLexer,
get_prompt_tokens=get_tokens,
reserve_space_for_menu=8,
)
buff = CypherBuffer(
accept_action=AcceptAction.RETURN_DOCUMENT,
history=FileHistory(filename=os.path.expanduser('~/.cycli_history')),
completer=completer,
complete_while_typing=True,
)
application = Application(
style=PygmentsStyle(CypherStyle),
buffer=buff,
layout=layout,
on_exit=AbortAction.RAISE_EXCEPTION,
key_bindings_registry=CypherBinder.registry
)
cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = cli.run()
query = document.text
self.handle_query(query)
except UserWantsOut:
print("Goodbye!")
except Exception as e:
print(e)
def run(self):
neo4j = Neo4j(self.host, self.port, self.username, self.password, self.ssl)
neo4j.connect()
self.neo4j = neo4j
try:
labels = neo4j.labels()
relationship_types = neo4j.relationship_types()
properties = neo4j.properties()
except Unauthorized:
print("Unauthorized. See cycli --help for authorization instructions.")
return
except SocketError:
print("Connection refused. Is Neo4j turned on?")
return
if self.filename:
queries = self.filename.read()
queries = queries.split(";")[:-1]
for query in queries:
query += ";"
query = query.strip()
print("> " + query)
self.handle_query(query)
print()
return
click.secho(" ______ __ __ ______ __ __ ", fg="red")
click.secho("/\ ___\ /\ \_\ \ /\ ___\ /\ \ /\ \ ", fg="yellow")
click.secho("\ \ \____ \ \____ \ \ \ \____ \ \ \____ \ \ \ ", fg="green")
click.secho(" \ \_____\ \/\_____\ \ \_____\ \ \_____\ \ \_\ ", fg="blue")
click.secho(" \/_____/ \/_____/ \/_____/ \/_____/ \/_/ ", fg="magenta")
print("\nVersion: {}".format(__version__))
print("Bug reports: https://github.com/nicolewhite/cycli/issues\n")
completer = CypherCompleter(labels, relationship_types, properties)
layout = create_default_layout(
lexer=CypherLexer,
get_prompt_tokens=get_tokens,
reserve_space_for_menu=True
)
buff = CypherBuffer(
history=FileHistory(filename=os.path.expanduser('~/.cycli_history')),
completer=completer,
complete_while_typing=Always()
)
application = Application(
style=CypherStyle,
buffer=buff,
layout=layout,
on_exit=AbortAction.RAISE_EXCEPTION,
key_bindings_registry=CypherBinder.registry
)
cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = cli.run()
query = document.text
self.handle_query(query)
except Exception:
print("Goodbye!")
class Saws(object):
"""Encapsulates the Saws CLI.
Attributes:
* aws_cli: An instance of prompt_toolkit's CommandLineInterface.
* config: An instance of Config.
* config_obj: An instance of ConfigObj, reads from ~/.sawsrc.
* aws_commands: An instance of AwsCommands
* commands: A list of commands from data/SOURCES.txt.
* sub_commands: A list of sub_commands from data/SOURCES.txt.
* global_options: A list of global_options from data/SOURCES.txt.
* resource_options: A list of resource_options from data/SOURCES.txt,
used for syntax coloring.
* ec2_states: A list of ec2_states from data/SOURCES.txt.
* completer: An instance of AwsCompleter.
* key_manager: An instance of KeyManager
* logger: An instance of SawsLogger.
* theme: A string representing the lexer theme.
Currently only 'vim' is supported.
"""
def __init__(self):
"""Inits Saws.
Args:
* None.
Returns:
None.
"""
self.aws_cli = None
self.key_manager = None
self.PYGMENTS_CMD = ' | pygmentize -l json'
self.config = Config()
self.config_obj = self.config.read_configuration()
self.theme = self.config_obj['main']['theme']
self.logger = SawsLogger(__name__,
self.config_obj['main']['log_file'],
self.config_obj['main']['log_level']).logger
self.aws_commands = AwsCommands()
self.commands, self.sub_commands, self.global_options, \
self.resource_options, self.ec2_states \
= self.aws_commands.generate_all_commands()
self.completer = AwsCompleter(
awscli_completer,
self.commands,
self.config_obj,
self.log_exception,
ec2_states=self.ec2_states,
fuzzy_match=self.get_fuzzy_match(),
shortcut_match=self.get_shortcut_match())
self.create_cli()
def log_exception(self, e, traceback, echo=False):
"""Logs the exception and traceback to the log file ~/.saws.log.
Args:
* e: A Exception that specifies the exception.
* traceback: A Traceback that specifies the traceback.
* echo: A boolean that specifies whether to echo the exception
to the console using click.
Returns:
None.
"""
self.logger.debug('exception: %r.', str(e))
self.logger.error("traceback: %r", traceback.format_exc())
if echo:
click.secho(str(e), fg='red')
def set_color(self, color):
"""Setter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the color flag.
Returns:
None.
"""
self.config_obj['main']['color_output'] = color
def get_color(self):
"""Getter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the color flag.
"""
return self.config_obj['main'].as_bool('color_output')
def set_fuzzy_match(self, fuzzy):
"""Setter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the fuzzy flag.
Returns:
None.
"""
self.config_obj['main']['fuzzy_match'] = fuzzy
self.completer.fuzzy_match = fuzzy
def get_fuzzy_match(self):
"""Getter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the fuzzy flag.
"""
return self.config_obj['main'].as_bool('fuzzy_match')
def set_shortcut_match(self, shortcut):
"""Setter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the shortcut flag.
Returns:
None.
"""
self.config_obj['main']['shortcut_match'] = shortcut
self.completer.shortcut_match = shortcut
def get_shortcut_match(self):
"""Getter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the shortcut flag.
"""
return self.config_obj['main'].as_bool('shortcut_match')
def refresh_resources(self):
"""Convenience function to refresh resources for completion.
Used by prompt_toolkit's KeyBindingManager.
Args:
* None.
Returns:
None.
"""
self.completer.refresh_resources(force_refresh=True)
def handle_docs(self, text=None, from_fkey=False):
"""Displays contextual web docs for `F9` or the `docs` command.
Displays the web docs specific to the currently entered:
* (optional) command
* (optional) subcommand
If no command or subcommand is present, the docs index page is shown.
Docs are only displayed if:
* from_fkey is True
* from_fkey is False and `docs` is found in text
Args:
* text: A string representing the input command text.
* from_fkey: A boolean representing whether this function is
being executed from an `F9` key press.
Returns:
A boolean representing whether the web docs were shown.
"""
base_url = 'http://docs.aws.amazon.com/cli/latest/reference/'
index_html = 'index.html'
if text is None:
text = self.aws_cli.current_buffer.document.text
# If the user hit the F9 key, append 'docs' to the text
if from_fkey:
text = text.strip() + ' ' + AwsCommands.AWS_DOCS
tokens = text.split()
if len(tokens) > 2 and tokens[-1] == AwsCommands.AWS_DOCS:
prev_word = tokens[-2]
# If we have a command, build the url
if prev_word in self.commands:
prev_word = prev_word + '/'
url = base_url + prev_word + index_html
webbrowser.open(url)
return True
# if we have a command and subcommand, build the url
elif prev_word in self.sub_commands:
command_url = tokens[-3] + '/'
sub_command_url = tokens[-2] + '.html'
url = base_url + command_url + sub_command_url
webbrowser.open(url)
return True
webbrowser.open(base_url + index_html)
# If we still haven't opened the help doc at this point and the
# user hit the F9 key or typed docs, just open the main docs index
if from_fkey or AwsCommands.AWS_DOCS in tokens:
webbrowser.open(base_url + index_html)
return True
return False
def handle_cd(self, text):
"""Handles a `cd` shell command by calling python's os.chdir.
Simply passing in the `cd` command to subprocess.call doesn't work.
Note: Changing the directory within Saws will only be in effect while
running Saws. Exiting the program will return you to the directory
you were in prior to running Saws.
Attributes:
* text: A string representing the input command text.
Returns:
A boolean representing a `cd` command was found and handled.
"""
CD_CMD = 'cd'
stripped_text = text.strip()
if stripped_text.startswith(CD_CMD):
directory = ''
if stripped_text == CD_CMD:
# Treat `cd` as a change to the root directory.
# os.path.expanduser does this in a cross platform manner.
directory = os.path.expanduser('~')
else:
tokens = text.split(CD_CMD + ' ')
directory = tokens[-1]
try:
os.chdir(directory)
except OSError as e:
self.log_exception(e, traceback, echo=True)
return True
return False
def colorize_output(self, text):
"""Highlights output with pygments.
Only highlights the output if all of the following conditions are True:
* The color option is enabled
* The text does not contain the `configure` command
* The text does not contain the `help` command, which already does
output highlighting
Args:
* text: A string that represents the input command text.
Returns:
A string that represents:
* The original command text if no highlighting was performed.
* The pygments highlighted command text otherwise.
"""
if not self.get_color():
return text
stripped_text = text.strip()
excludes = [AwsCommands.AWS_CONFIGURE,
AwsCommands.AWS_HELP,
'|']
if not any(substring in stripped_text for substring in excludes):
return text.strip() + self.PYGMENTS_CMD
else:
return text
def process_command(self, text):
"""Processes the input command, called by the cli event loop
Args:
* text: A string that represents the input command text.
Returns:
None.
"""
if AwsCommands.AWS_COMMAND in text:
text = self.completer.replace_shortcut(text)
if self.handle_docs(text):
return
text = self.colorize_output(text)
try:
if not self.handle_cd(text):
# Pass the command onto the shell so aws-cli can execute it
subprocess.call(text, shell=True)
print('')
except Exception as e:
self.log_exception(e, traceback, echo=True)
def create_cli(self):
"""Creates the prompt_toolkit's CommandLineInterface.
Long description.
Args:
* None.
Returns:
None.
"""
history = FileHistory(os.path.expanduser('~/.saws-history'))
toolbar = Toolbar(self.get_color,
self.get_fuzzy_match,
self.get_shortcut_match)
layout = create_default_layout(
message='saws> ',
reserve_space_for_menu=True,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar.handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
completer=self.completer,
complete_while_typing=Always())
self.key_manager = KeyManager(
self.set_color,
self.get_color,
self.set_fuzzy_match,
self.get_fuzzy_match,
self.set_shortcut_match,
self.get_shortcut_match,
self.refresh_resources,
self.handle_docs)
style_factory = StyleFactory(self.theme)
application = Application(
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
ignore_case=True)
eventloop = create_eventloop()
self.aws_cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def run_cli(self):
"""Runs the main loop.
Args:
* None.
Returns:
None.
"""
print('Version:', __version__)
print('Theme:', self.theme)
while True:
document = self.aws_cli.run()
self.process_command(document.text)
class WharfeeCli(object):
"""
The CLI implementation.
"""
dcli = None
keyword_completer = None
handler = None
saved_less_opts = None
config = None
config_template = 'wharfeerc'
config_name = '~/.wharfeerc'
def __init__(self, no_completion=False):
"""
Initialize class members.
Should read the config here at some point.
"""
self.config = self.read_configuration()
self.theme = self.config['main']['theme']
log_file = self.config['main']['log_file']
log_level = self.config['main']['log_level']
self.logger = create_logger(__name__, log_file, log_level)
# set_completer_options refreshes all by default
self.handler = DockerClient(
self.config['main'].as_int('client_timeout'),
self.clear,
self.refresh_completions_force,
self.logger)
self.completer = DockerCompleter(
long_option_names=self.get_long_options(),
fuzzy=self.get_fuzzy_match())
self.set_completer_options()
self.completer.set_enabled(not no_completion)
self.saved_less_opts = self.set_less_opts()
def read_configuration(self):
"""
:return:
"""
default_config = os.path.join(
self.get_package_path(), self.config_template)
write_default_config(default_config, self.config_name)
return read_config(self.config_name, default_config)
def get_package_path(self):
"""
Find out pakage root path.
:return: string: path
"""
from wharfee import __file__ as package_root
return os.path.dirname(package_root)
def set_less_opts(self):
"""
Set the "less" options and save the old settings.
What we're setting:
-F:
--quit-if-one-screen: Quit if entire file fits on first screen.
-R:
--raw-control-chars: Output "raw" control characters.
-X:
--no-init: Don't use termcap keypad init/deinit strings.
--no-keypad: Don't use termcap init/deinit strings.
This also disables launching "less" in an alternate screen.
:return: string with old options
"""
opts = os.environ.get('LESS', '')
os.environ['LESS'] = '-RXF'
return opts
def revert_less_opts(self):
"""
Restore the previous "less" options.
"""
os.environ['LESS'] = self.saved_less_opts
def write_config_file(self):
"""
Write config file on exit.
"""
self.config.write()
def clear(self):
"""
Clear the screen.
"""
click.clear()
def set_completer_options(self, cons=True, runs=True, imgs=True, vols=True):
"""
Set image and container names in Completer.
Re-read if needed after a command.
:param cons: boolean: need to refresh containers
:param runs: boolean: need to refresh running containers
:param imgs: boolean: need to refresh images
:param vols: boolean: need to refresh volumes
"""
if cons:
cs = self.handler.containers(all=True)
if cs and len(cs) > 0 and isinstance(cs[0], dict):
containers = [name for c in cs for name in c['Names']]
self.completer.set_containers(containers)
if runs:
cs = self.handler.containers()
if cs and len(cs) > 0 and isinstance(cs[0], dict):
running = [name for c in cs for name in c['Names']]
self.completer.set_running(running)
if imgs:
def format_tagged(tagname, img_id):
if tagname == '<none>:<none>':
return img_id[:11]
return tagname
def parse_image_name(tag, img_id):
if ':' in tag:
result = tag.split(':', 2)[0]
else:
result = tag
if result == '<none>':
result = img_id[:11]
return result
ims = self.handler.images()
if ims and len(ims) > 0 and isinstance(ims[0], dict):
images = set([])
tagged = set([])
for im in ims:
repo_tag = '{0}:{1}'.format(im['Repository'], im['Tag'])
images.add(parse_image_name(repo_tag, im['Id']))
tagged.add(format_tagged(repo_tag, im['Id']))
self.completer.set_images(images)
self.completer.set_tagged(tagged)
if vols:
vs = self.handler.volume_ls(quiet=True)
self.completer.set_volumes(vs)
def set_fuzzy_match(self, is_fuzzy):
"""
Setter for fuzzy matching mode
:param is_fuzzy: boolean
"""
self.config['main']['fuzzy_match'] = is_fuzzy
self.completer.set_fuzzy_match(is_fuzzy)
def get_fuzzy_match(self):
"""
Getter for fuzzy matching mode
:return: boolean
"""
return self.config['main'].as_bool('fuzzy_match')
def set_long_options(self, is_long):
"""
Setter for long option names.
:param is_long: boolean
"""
self.config['main']['suggest_long_option_names'] = is_long
self.completer.set_long_options(is_long)
def get_long_options(self):
"""
Getter for long option names.
:return: boolean
"""
return self.config['main'].as_bool('suggest_long_option_names')
def refresh_completions_force(self):
"""Force refresh and make it visible."""
self.set_completer_options()
click.echo('Refreshed completions.')
def refresh_completions(self):
"""
After processing the command, refresh the lists of
containers and images as needed
"""
self.set_completer_options(self.handler.is_refresh_containers,
self.handler.is_refresh_running,
self.handler.is_refresh_images,
self.handler.is_refresh_volumes)
def run_cli(self):
"""
Run the main loop
"""
print('Version:', __version__)
print('Home: http://wharfee.com')
history = FileHistory(os.path.expanduser('~/.wharfee-history'))
toolbar_handler = create_toolbar_handler(self.get_long_options, self.get_fuzzy_match)
layout = create_prompt_layout(
message='wharfee> ',
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar_handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
manager = get_key_manager(
self.set_long_options,
self.get_long_options,
self.set_fuzzy_match,
self.get_fuzzy_match)
application = Application(
style=style_factory(self.theme),
layout=layout,
buffer=cli_buffer,
key_bindings_registry=manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.dcli = CommandLineInterface(
application=application,
eventloop=eventloop)
while True:
try:
document = self.dcli.run(True)
self.handler.handle_input(document.text)
if isinstance(self.handler.output, GeneratorType):
output_stream(self.handler.command,
self.handler.output,
self.handler.log)
elif self.handler.output is not None:
lines = format_data(
self.handler.command,
self.handler.output)
click.echo_via_pager('\n'.join(lines))
if self.handler.after:
for line in self.handler.after():
click.echo(line)
if self.handler.exception:
# This was handled, just log it.
self.logger.warning('An error was handled: %r',
self.handler.exception)
self.refresh_completions()
except OptionError as ex:
self.logger.debug('Error: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(ex.msg, fg='red')
except KeyboardInterrupt:
# user pressed Ctrl + C
if self.handler.after:
click.echo('')
for line in self.handler.after():
click.echo(line)
self.refresh_completions()
except DockerPermissionException as ex:
self.logger.debug('Permission exception: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(str(ex), fg='red')
except EOFError:
# exit out of the CLI
break
# TODO: uncomment for release
except Exception as ex:
self.logger.debug('Exception: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(str(ex), fg='red')
self.revert_less_opts()
self.write_config_file()
print('Goodbye!')
class WharfeeCli(object):
"""
The CLI implementation.
"""
dcli = None
keyword_completer = None
handler = None
saved_less_opts = None
config = None
config_template = "wharfeerc"
config_name = "~/.wharfeerc"
def __init__(self):
"""
Initialize class members.
Should read the config here at some point.
"""
self.config = self.read_configuration()
self.theme = self.config["main"]["theme"]
log_file = self.config["main"]["log_file"]
log_level = self.config["main"]["log_level"]
self.logger = create_logger(__name__, log_file, log_level)
# set_completer_options refreshes all by default
self.handler = DockerClient(
self.config["main"].as_int("client_timeout"), self.clear, self.set_completer_options
)
self.completer = DockerCompleter(long_option_names=self.get_long_options(), fuzzy=self.get_fuzzy_match())
self.set_completer_options()
self.saved_less_opts = self.set_less_opts()
def read_configuration(self):
"""
:return:
"""
default_config = os.path.join(self.get_package_path(), self.config_template)
write_default_config(default_config, self.config_name)
return read_config(self.config_name, default_config)
def get_package_path(self):
"""
Find out pakage root path.
:return: string: path
"""
from wharfee import __file__ as package_root
return os.path.dirname(package_root)
def set_less_opts(self):
"""
Set the "less" options and save the old settings.
What we're setting:
-F:
--quit-if-one-screen: Quit if entire file fits on first screen.
-R:
--raw-control-chars: Output "raw" control characters.
-X:
--no-init: Don't use termcap keypad init/deinit strings.
--no-keypad: Don't use termcap init/deinit strings.
This also disables launching "less" in an alternate screen.
:return: string with old options
"""
opts = os.environ.get("LESS", "")
os.environ["LESS"] = "-RXF"
return opts
def revert_less_opts(self):
"""
Restore the previous "less" options.
"""
os.environ["LESS"] = self.saved_less_opts
def write_config_file(self):
"""
Write config file on exit.
"""
self.config.write()
def clear(self):
"""
Clear the screen.
"""
click.clear()
def set_completer_options(self, cons=True, runs=True, imgs=True):
"""
Set image and container names in Completer.
Re-read if needed after a command.
:param cons: boolean: need to refresh containers
:param runs: boolean: need to refresh running containers
:param imgs: boolean: need to refresh images
"""
if cons:
cs = self.handler.containers(all=True)
if cs and len(cs) > 0 and isinstance(cs[0], dict):
containers = [name for c in cs for name in c["Names"]]
self.completer.set_containers(containers)
if runs:
cs = self.handler.containers()
if cs and len(cs) > 0 and isinstance(cs[0], dict):
running = [name for c in cs for name in c["Names"]]
self.completer.set_running(running)
if imgs:
def format_tagged(tagname, img_id):
if tagname == "<none>:<none>":
return img_id[:11]
return tagname
def parse_image_name(tag, img_id):
if ":" in tag:
result = tag.split(":", 2)[0]
else:
result = tag
if result == "<none>":
result = img_id[:11]
return result
ims = self.handler.images()
if ims and len(ims) > 0 and isinstance(ims[0], dict):
images = set([])
tagged = set([])
for im in ims:
repo_tag = "{0}:{1}".format(im["Repository"], im["Tag"])
images.add(parse_image_name(repo_tag, im["Id"]))
tagged.add(format_tagged(repo_tag, im["Id"]))
self.completer.set_images(images)
self.completer.set_tagged(tagged)
def set_fuzzy_match(self, is_fuzzy):
"""
Setter for fuzzy matching mode
:param is_fuzzy: boolean
"""
self.config["main"]["fuzzy_match"] = is_fuzzy
self.completer.set_fuzzy_match(is_fuzzy)
def get_fuzzy_match(self):
"""
Getter for fuzzy matching mode
:return: boolean
"""
return self.config["main"].as_bool("fuzzy_match")
def set_long_options(self, is_long):
"""
Setter for long option names.
:param is_long: boolean
"""
self.config["main"]["suggest_long_option_names"] = is_long
self.completer.set_long_options(is_long)
def get_long_options(self):
"""
Getter for long option names.
:return: boolean
"""
return self.config["main"].as_bool("suggest_long_option_names")
def refresh_completions(self):
"""
After processing the command, refresh the lists of
containers and images as needed
"""
self.set_completer_options(
self.handler.is_refresh_containers, self.handler.is_refresh_running, self.handler.is_refresh_images
)
def run_cli(self):
"""
Run the main loop
"""
print("Version:", __version__)
print("Home: http://wharfee.com")
history = FileHistory(os.path.expanduser("~/.wharfee-history"))
toolbar_handler = create_toolbar_handler(self.get_long_options, self.get_fuzzy_match)
layout = create_default_layout(
message="wharfee> ",
reserve_space_for_menu=True,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar_handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars="[](){}"),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone(),
)
],
)
cli_buffer = Buffer(history=history, completer=self.completer, complete_while_typing=Always())
manager = get_key_manager(
self.set_long_options, self.get_long_options, self.set_fuzzy_match, self.get_fuzzy_match
)
application = Application(
style=style_factory(self.theme),
layout=layout,
buffer=cli_buffer,
key_bindings_registry=manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
ignore_case=True,
)
eventloop = create_eventloop()
self.dcli = CommandLineInterface(application=application, eventloop=eventloop)
while True:
try:
document = self.dcli.run()
self.handler.handle_input(document.text)
if isinstance(self.handler.output, GeneratorType):
output_stream(self.handler.command, self.handler.output, self.handler.logs)
elif self.handler.output is not None:
lines = format_data(self.handler.command, self.handler.output)
click.echo_via_pager("\n".join(lines))
if self.handler.after:
for line in self.handler.after():
click.echo(line)
self.refresh_completions()
except OptionError as ex:
self.logger.debug("Error: %r.", ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(ex.msg, fg="red")
except KeyboardInterrupt:
# user pressed Ctrl + C
if self.handler.after:
click.echo("")
for line in self.handler.after():
click.echo(line)
self.refresh_completions()
except DockerPermissionException as ex:
self.logger.debug("Permission exception: %r.", ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(ex.message, fg="red")
except EOFError:
# exit out of the CLI
break
# TODO: uncomment for release
except Exception as ex:
self.logger.debug("Exception: %r.", ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho("{0}".format(ex), fg="red")
break
self.revert_less_opts()
self.write_config_file()
print("Goodbye!")
class Prompt:
def __init__(self, run_name, state_obj):
self.run_name = run_name
self.state_obj = state_obj
self.cli = None
self.q = queue.Queue()
self.thread = threading.Thread(target=self.run)
def start(self):
self.thread.start()
def stop(self):
if self.cli:
self.cli.exit()
self.thread.join()
def get_bottom_toolbar_tokens(self, cli):
return [(Token.Toolbar, 'Run '), (Token.Name, self.run_name),
(Token.Toolbar, ' in progress.')]
def get_prompt_tokens(self, cli):
return [(Token.Prompt, '> ')]
def run(self):
style = style_from_dict({
Token.Prompt: 'bold',
Token.Toolbar: '#ccc bg:#333',
Token.Name: '#fff bold bg:#333',
})
history = InMemoryHistory()
eventloop = create_eventloop()
app = create_prompt_application(
history=history,
style=style,
get_bottom_toolbar_tokens=self.get_bottom_toolbar_tokens,
get_prompt_tokens=self.get_prompt_tokens)
self.cli = CommandLineInterface(app, eventloop)
with self.cli.patch_stdout_context(raw=True):
while True:
try:
self.cli.run()
doc = self.cli.return_value()
if doc is None:
return
cmd = shlex.split(doc.text)
app.buffer.reset(append_to_history=True)
if not cmd:
continue
elif cmd[0] in ('exit', 'quit'):
self.q.put(Exit())
return
elif cmd[0] == 'help':
print('Help text forthcoming.')
elif cmd[0] == 'skip':
self.q.put(Skip())
elif cmd[0] == 'set':
self.q.put(
Set(cmd[1], ast.literal_eval(' '.join(cmd[2:]))))
else:
print('Unknown command. Try \'help\'.')
except KeyboardInterrupt:
continue
except EOFError:
self.q.put(Exit())
return
except Exception as err:
print(err)
self.q.put(Exit())
return
class MyCli(object):
default_prompt = "\\t \\u@\\h:\\d> "
defaults_suffix = None
# In order of being loaded. Files lower in list override earlier ones.
cnf_files = ["/etc/my.cnf", "/etc/mysql/my.cnf", "/usr/local/etc/my.cnf", "~/.my.cnf"]
system_config_files = ["/etc/myclirc"]
default_config_file = os.path.join(PACKAGE_ROOT, "myclirc")
user_config_file = "~/.myclirc"
def __init__(
self,
sqlexecute=None,
prompt=None,
logfile=None,
defaults_suffix=None,
defaults_file=None,
login_path=None,
auto_vertical_output=False,
warn=None,
):
self.sqlexecute = sqlexecute
self.logfile = logfile
self.defaults_suffix = defaults_suffix
self.login_path = login_path
self.auto_vertical_output = auto_vertical_output
# self.cnf_files is a class variable that stores the list of mysql
# config files to read in at launch.
# If defaults_file is specified then override the class variable with
# defaults_file.
if defaults_file:
self.cnf_files = [defaults_file]
# Load config.
config_files = [self.default_config_file] + self.system_config_files + [self.user_config_file]
c = self.config = read_config_files(config_files)
self.multi_line = c["main"].as_bool("multi_line")
self.key_bindings = c["main"]["key_bindings"]
special.set_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")
c_dest_warning = c["main"].as_bool("destructive_warning")
self.destructive_warning = c_dest_warning if warn is None else warn
# Write user config if system config wasn't the last config loaded.
if c.filename not in self.system_config_files:
write_default_config(self.default_config_file, self.user_config_file)
# audit log
if self.logfile is None and "audit_log" in c["main"]:
try:
self.logfile = open(os.path.expanduser(c["main"]["audit_log"]), "a")
except (IOError, OSError) as e:
self.output(
"Error: Unable to open the audit log file. Your queries will not be logged.", err=True, fg="red"
)
self.logfile = False
self.completion_refresher = CompletionRefresher()
self.logger = logging.getLogger(__name__)
self.initialize_logging()
prompt_cnf = self.read_my_cnf_files(self.cnf_files, ["prompt"])["prompt"]
self.prompt_format = prompt or prompt_cnf or c["main"]["prompt"] or self.default_prompt
self.query_history = []
# Initialize completer.
smart_completion = c["main"].as_bool("smart_completion")
self.completer = SQLCompleter(smart_completion)
self._completer_lock = threading.Lock()
# Register custom special commands.
self.register_special_commands()
# Load .mylogin.cnf if it exists.
mylogin_cnf_path = get_mylogin_cnf_path()
if mylogin_cnf_path:
try:
mylogin_cnf = open_mylogin_cnf(mylogin_cnf_path)
if mylogin_cnf_path and mylogin_cnf:
# .mylogin.cnf gets read last, even if defaults_file is specified.
self.cnf_files.append(mylogin_cnf)
elif mylogin_cnf_path and not mylogin_cnf:
# There was an error reading the login path file.
print("Error: Unable to read login path file.")
except CryptoError:
click.secho("Warning: .mylogin.cnf was not read: pycrypto " "module is not available.")
self.cli = None
def register_special_commands(self):
special.register_special_command(self.change_db, "use", "\\u", "Change to a new database.", aliases=("\\u",))
special.register_special_command(
self.change_db,
"connect",
"\\r",
"Reconnect to the database. Optional database argument.",
aliases=("\\r",),
case_sensitive=True,
)
special.register_special_command(
self.refresh_completions, "rehash", "\\#", "Refresh auto-completions.", arg_type=NO_QUERY, aliases=("\\#",)
)
special.register_special_command(
self.change_table_format, "tableformat", "\\T", "Change Table Type.", aliases=("\\T",), case_sensitive=True
)
special.register_special_command(
self.execute_from_file, "source", "\\. filename", "Execute commands from file.", aliases=("\\.",)
)
special.register_special_command(
self.change_prompt_format, "prompt", "\\R", "Change prompt format.", aliases=("\\R",), case_sensitive=True
)
def change_table_format(self, arg, **_):
if not arg in table_formats():
msg = "Table type %s not yet implemented. Allowed types:" % arg
for table_type in table_formats():
msg += "\n\t%s" % table_type
yield (None, None, None, msg)
else:
self.table_format = arg
yield (None, None, None, "Changed table Type to %s" % self.table_format)
def change_db(self, arg, **_):
if arg is None:
self.sqlexecute.connect()
else:
self.sqlexecute.connect(database=arg)
yield (
None,
None,
None,
'You are now connected to database "%s" as ' 'user "%s"' % (self.sqlexecute.dbname, self.sqlexecute.user),
)
def execute_from_file(self, arg, **_):
if not arg:
message = "Missing required argument, filename."
return [(None, None, None, message)]
try:
with open(os.path.expanduser(arg), encoding="utf-8") as f:
query = f.read()
except IOError as e:
return [(None, None, None, str(e))]
if self.destructive_warning and confirm_destructive_query(query) is False:
message = "Wise choice. Command execution stopped."
return [(None, None, None, message)]
return self.sqlexecute.run(query)
def change_prompt_format(self, arg, **_):
"""
Change the prompt format.
"""
if not arg:
message = "Missing required argument, format."
return [(None, None, None, message)]
self.prompt_format = self.get_prompt(arg)
return [(None, None, None, "Changed prompt format to %s" % arg)]
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("mycli")
root_logger.addHandler(handler)
root_logger.setLevel(level_map[log_level.upper()])
# Only capture warnings on Python 2.7 and later.
try:
logging.captureWarnings(True)
except AttributeError:
pass
root_logger.debug("Initializing mycli logging.")
root_logger.debug("Log file %r.", log_file)
def connect_uri(self, uri, local_infile=None, ssl=None):
uri = urlparse(uri)
database = uri.path[1:] # ignore the leading fwd slash
self.connect(database, uri.username, uri.password, uri.hostname, uri.port, local_infile=local_infile, ssl=ssl)
def read_my_cnf_files(self, files, keys):
"""
Reads a list of config files and merges them. The last one will win.
:param files: list of files to read
:param keys: list of keys to retrieve
:returns: tuple, with None for missing keys.
"""
cnf = read_config_files(files)
sections = ["client"]
if self.login_path and self.login_path != "client":
sections.append(self.login_path)
if self.defaults_suffix:
sections.extend([sect + self.defaults_suffix for sect in sections])
def get(key):
result = None
for sect in cnf:
if sect in sections and key in cnf[sect]:
result = cnf[sect][key]
return result
return dict([(x, get(x)) for x in keys])
def merge_ssl_with_cnf(self, ssl, cnf):
"""Merge SSL configuration dict with cnf dict"""
merged = {}
merged.update(ssl)
prefix = "ssl-"
for k, v in cnf.items():
# skip unrelated options
if not k.startswith(prefix):
continue
if v is None:
continue
# special case because PyMySQL argument is significantly different
# from commandline
if k == "ssl-verify-server-cert":
merged["check_hostname"] = v
else:
# use argument name just strip "ssl-" prefix
arg = k[len(prefix) :]
merged[arg] = v
return merged
def connect(
self, database="", user="", passwd="", host="", port="", socket="", charset="", local_infile="", ssl=""
):
cnf = {
"database": None,
"user": None,
"password": None,
"host": None,
"port": None,
"socket": None,
"default-character-set": None,
"local-infile": None,
"loose-local-infile": None,
"ssl-ca": None,
"ssl-cert": None,
"ssl-key": None,
"ssl-cipher": None,
"ssl-verify-serer-cert": None,
}
cnf = self.read_my_cnf_files(self.cnf_files, cnf.keys())
# Fall back to config values only if user did not specify a value.
database = database or cnf["database"]
if port or host:
socket = ""
else:
socket = socket or cnf["socket"]
user = user or cnf["user"] or os.getenv("USER")
host = host or cnf["host"] or "localhost"
port = port or cnf["port"] or 3306
ssl = ssl or {}
try:
port = int(port)
except ValueError as e:
self.output("Error: Invalid port number: '{0}'.".format(port), err=True, fg="red")
exit(1)
passwd = passwd or cnf["password"]
charset = charset or cnf["default-character-set"] or "utf8"
# Favor whichever local_infile option is set.
for local_infile_option in (local_infile, cnf["local-infile"], cnf["loose-local-infile"], False):
try:
local_infile = str_to_bool(local_infile_option)
break
except (TypeError, ValueError):
pass
ssl = self.merge_ssl_with_cnf(ssl, cnf)
# prune lone check_hostname=False
if not any(v for v in ssl.values()):
ssl = None
# Connect to the database.
try:
try:
sqlexecute = SQLExecute(database, user, passwd, host, port, socket, charset, local_infile, ssl)
except OperationalError as e:
if "Access denied for user" in e.args[1]:
passwd = click.prompt("Password", hide_input=True, show_default=False, type=str)
sqlexecute = SQLExecute(database, user, passwd, host, port, socket, charset, local_infile, ssl)
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())
self.output(str(e), err=True, fg="red")
exit(1)
self.sqlexecute = sqlexecute
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):
sqlexecute = self.sqlexecute
logger = self.logger
self.configure_pager()
self.refresh_completions()
project_root = os.path.dirname(PACKAGE_ROOT)
author_file = os.path.join(project_root, "AUTHORS")
sponsor_file = os.path.join(project_root, "SPONSORS")
key_binding_manager = mycli_bindings()
print("Version:", __version__)
print("Chat: https://gitter.im/dbcli/mycli")
print("Mail: https://groups.google.com/forum/#!forum/mycli-users")
print("Home: http://mycli.net")
print("Thanks to the contributor -", thanks_picker([author_file, sponsor_file]))
def prompt_tokens(cli):
return [(Token.Prompt, self.get_prompt(self.prompt_format))]
def get_continuation_tokens(cli, width):
return [(Token.Continuation, " " * (width - 3) + "-> ")]
get_toolbar_tokens = create_toolbar_tokens_func(self.completion_refresher.is_refreshing)
layout = create_prompt_layout(
lexer=MyCliLexer,
multiline=True,
get_prompt_tokens=prompt_tokens,
get_continuation_tokens=get_continuation_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
display_completions_in_columns=self.wider_completion_menu,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars="[](){}"),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone(),
)
],
)
with self._completer_lock:
buf = CLIBuffer(
always_multiline=self.multi_line,
completer=self.completer,
history=FileHistory(os.path.expanduser("~/.mycli-history")),
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT,
)
if self.key_bindings == "vi":
editing_mode = EditingMode.VI
else:
editing_mode = EditingMode.EMACS
application = Application(
style=style_factory(self.syntax_style, self.cli_style),
layout=layout,
buffer=buf,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
editing_mode=editing_mode,
ignore_case=True,
)
self.cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = self.cli.run(reset_current_buffer=True)
special.set_expanded_output(False)
# The reason we check here instead of inside the sqlexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the
# sqlexecute.run() statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(self.cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg="red")
continue
if self.destructive_warning:
destroy = confirm_destructive_query(document.text)
if destroy is None:
pass # Query was not destructive. Nothing to do here.
elif destroy is True:
self.output("Your call!")
else:
self.output("Wise choice!")
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug("sql: %r", document.text)
if self.logfile:
self.logfile.write("\n# %s\n" % datetime.now())
self.logfile.write(document.text)
self.logfile.write("\n")
successful = False
start = time()
res = sqlexecute.run(document.text)
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
threshold = 1000
if is_select(status) and cur and cur.rowcount > threshold:
self.output("The result set has more than %s rows." % threshold, fg="red")
if not click.confirm("Do you want to continue?"):
self.output("Aborted!", err=True, fg="red")
break
if self.auto_vertical_output:
max_width = self.cli.output.get_size().columns
else:
max_width = None
formatted = format_output(
title, cur, headers, status, self.table_format, special.is_expanded_output(), max_width
)
output.extend(formatted)
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except UnicodeDecodeError as e:
import pymysql
if pymysql.VERSION < (0, 6, 7):
message = (
"You are running an older version of pymysql.\n"
"Please upgrade to 0.6.7 or above to view binary data.\n"
"Try 'pip install -U pymysql'."
)
self.output(message)
else:
raise e
except KeyboardInterrupt:
# Restart connection to the database
sqlexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
self.output("cancelled query", err=True, fg="red")
except NotImplementedError:
self.output("Not Yet Implemented.", fg="yellow")
except OperationalError as e:
logger.debug("Exception: %r", e)
reconnect = True
if e.args[0] in (2003, 2006, 2013):
reconnect = click.prompt(
"Connection reset. Reconnect (Y/n)", show_default=False, type=bool, default=True
)
if reconnect:
logger.debug("Attempting to reconnect.")
try:
sqlexecute.connect()
logger.debug("Reconnected successfully.")
self.output("Reconnected!\nTry the command again.", fg="green")
except OperationalError as e:
logger.debug("Reconnect failed. e: %r", e)
self.output(str(e), err=True, fg="red")
continue # If reconnection failed, don't proceed further.
else: # If user chooses not to reconnect, don't proceed further.
continue
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg="red")
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg="red")
else:
try:
if special.is_pager_enabled():
self.output_via_pager("\n".join(output))
else:
self.output("\n".join(output))
except KeyboardInterrupt:
pass
if special.is_timing_enabled():
self.output("Time: %0.03fs" % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions(reset=need_completion_reset(document.text))
finally:
if self.logfile is False:
self.output("Warning: This query was not logged.", err=True, fg="red")
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
self.output("Goodbye!")
def output(self, text, **kwargs):
if self.logfile:
self.logfile.write(utf8tounicode(text))
self.logfile.write("\n")
click.secho(text, **kwargs)
def output_via_pager(self, text):
if self.logfile:
self.logfile.write(text)
self.logfile.write("\n")
click.echo_via_pager(text)
def configure_pager(self):
# Provide sane defaults for less.
os.environ["LESS"] = "-RXF"
cnf = self.read_my_cnf_files(self.cnf_files, ["pager", "skip-pager"])
if cnf["pager"]:
special.set_pager(cnf["pager"])
if cnf["skip-pager"]:
special.disable_pager()
def refresh_completions(self, reset=False):
if reset:
with self._completer_lock:
self.completer.reset_completions()
self.completion_refresher.refresh(self.sqlexecute, 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 mycli 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):
sqlexecute = self.sqlexecute
string = string.replace("\\u", sqlexecute.user or "(none)")
string = string.replace("\\h", sqlexecute.host or "(none)")
string = string.replace("\\d", sqlexecute.dbname or "(none)")
string = string.replace("\\t", sqlexecute.server_type()[0] or "mycli")
string = string.replace("\\n", "\n")
return string
class VCli(object):
def __init__(self, vexecute=None, vclirc_file=None):
self.vexecute = vexecute
from vcli import __file__ as package_root
package_root = os.path.dirname(package_root)
default_config = os.path.join(package_root, 'vclirc')
write_default_config(default_config, vclirc_file)
self.vspecial = VSpecial()
# Load config.
c = self.config = load_config(vclirc_file, default_config)
self.multi_line = c['main'].as_bool('multi_line')
self.vi_mode = c['main'].as_bool('vi')
self.vspecial.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 = VCompleter(smart_completion, vspecial=self.vspecial)
self.completer = completer
self._completer_lock = threading.Lock()
self.register_special_commands()
self.cli = None
def register_special_commands(self):
self.vspecial.register(self.change_db,
'\\c',
'\\c[onnect] [DBNAME]',
'Connect to a new database',
aliases=('use', '\\connect', 'USE'))
self.vspecial.register(self.refresh_completions,
'\\#',
'\\#',
'Refresh auto-completions',
arg_type=NO_QUERY)
self.vspecial.register(self.refresh_completions,
'\\refresh',
'\\refresh',
'Refresh auto-completions',
arg_type=NO_QUERY)
def change_db(self, pattern, **_):
if pattern:
db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern
self.vexecute.connect(database=db)
else:
self.vexecute.connect()
yield (None, None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.vexecute.dbname, self.vexecute.user), True)
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('vcli')
root_logger.addHandler(handler)
root_logger.setLevel(level_map[log_level.upper()])
root_logger.debug('Initializing vcli logging.')
root_logger.debug('Log file %r.', log_file)
def connect_uri(self, uri):
uri = urlparse(uri)
database = uri.path[1:] # ignore the leading fwd slash
host = uri.hostname or 'localhost'
user = uri.username or getpass.getuser()
port = uri.port or 5433
password = uri.password or ''
self.connect(database, host, user, port, password)
def connect(self, database, host, user, port, passwd):
# Connect to the database
try:
self.vexecute = VExecute(database, user, passwd, host, port)
except errors.DatabaseError as e: # Connection can fail
self.logger.debug('Database connection failed: %r.', e)
self.logger.error("traceback: %r", traceback.format_exc())
error_msg = str(e) or type(e).__name__
click.secho(error_msg, err=True, fg='red')
exit(1)
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):
vexecute = self.vexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
completer = self.completer
self.refresh_completions()
def set_vi_mode(value):
self.vi_mode = value
key_binding_manager = vcli_bindings(
get_vi_mode_enabled=lambda: self.vi_mode,
set_vi_mode_enabled=set_vi_mode)
click.secho('Version: %s' % __version__)
def prompt_tokens(cli):
return [(Token.Prompt, '%s=> ' % vexecute.dbname)]
get_toolbar_tokens = create_toolbar_tokens_func(
lambda: self.vi_mode, self.completion_refresher.is_refreshing)
input_processors = [
# Highlight matching brackets while editing.
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
layout = create_prompt_layout(
lexer=PostgresLexer,
reserve_space_for_menu=8,
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=input_processors)
history_file = self.config['main']['history_file']
with self._completer_lock:
buf = VBuffer(always_multiline=self.multi_line,
completer=self.completer,
history=FileHistory(
os.path.expanduser(history_file)),
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)
self.cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = self.cli.run()
# The reason we check here instead of inside the vexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the vexecute.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 vexecute.run(). Which causes
# finally clause to fail.
res = []
start = time()
# Run the query.
res = vexecute.run(document.text, self.vspecial)
file_output = None
stdout_output = []
for title, cur, headers, status, force_stdout 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
formatted = format_output(
title, cur, headers, status, self.table_format,
self.vspecial.expanded_output,
self.vspecial.aligned, self.vspecial.show_header)
if self.vspecial.output is not sys.stdout:
file_output = self.vspecial.output
if force_stdout or not file_output:
output = stdout_output
else:
output = file_output
write_output(output, formatted)
if hasattr(cur, 'rowcount'):
if self.vspecial.show_header:
if cur.rowcount == 1:
write_output(output, '(1 row)')
elif headers:
rowcount = max(cur.rowcount, 0)
write_output(output,
'(%d rows)' % rowcount)
if document.text.startswith(
'\\') and cur.rowcount == 0:
stdout_output = [
'No matching relations found.'
]
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
vexecute.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 errors.ConnectionError as e:
reconnect = click.prompt(
'Connection reset. Reconnect (Y/n)',
show_default=False,
type=bool,
default=True)
if reconnect:
try:
vexecute.connect()
click.secho('Reconnected!\nTry the command again.',
fg='green')
except errors.DatabaseError as e:
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(format_error(e), err=True, fg='red')
else:
successful = True
if stdout_output:
output = '\n'.join(stdout_output)
try:
click.echo_via_pager(output)
except KeyboardInterrupt:
pass
if file_output:
try:
file_output.flush()
except KeyboardInterrupt:
pass
if self.vspecial.timing_enabled:
print('Time: %0.03fs' % (time() - start))
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions(
need_completion_reset(document.text))
# 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(
vexecute.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'] = '-RXF'
return less_opts
def refresh_completions(self, reset=False):
if reset:
with self._completer_lock:
self.completer.reset_completions()
self.completion_refresher.refresh(self.vexecute, self.vspecial,
self._on_completions_refreshed)
return [(None, None, None,
'Auto-completion refresh started in the background.', True)]
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)
class Saws(object):
"""Encapsulates the Saws CLI.
Attributes:
* aws_cli: An instance of prompt_toolkit's CommandLineInterface.
* key_manager: An instance of KeyManager.
* config: An instance of Config.
* config_obj: An instance of ConfigObj, reads from ~/.sawsrc.
* theme: A string representing the lexer theme.
* logger: An instance of SawsLogger.
* all_commands: A list of all commands, sub_commands, options, etc
from data/SOURCES.txt.
* commands: A list of commands from data/SOURCES.txt.
* sub_commands: A list of sub_commands from data/SOURCES.txt.
* completer: An instance of AwsCompleter.
"""
PYGMENTS_CMD = ' | pygmentize -l json'
def __init__(self, refresh_resources=True):
"""Inits Saws.
Args:
* refresh_resources: A boolean that determines whether to
refresh resources.
Returns:
None.
"""
self.aws_cli = None
self.key_manager = None
self.config = Config()
self.config_obj = self.config.read_configuration()
self.theme = self.config_obj[self.config.MAIN][self.config.THEME]
self.logger = SawsLogger(
__name__,
self.config_obj[self.config.MAIN][self.config.LOG_FILE],
self.config_obj[self.config.MAIN][self.config.LOG_LEVEL]).logger
self.all_commands = AwsCommands().all_commands
self.commands = \
self.all_commands[AwsCommands.CommandType.COMMANDS.value]
self.sub_commands = \
self.all_commands[AwsCommands.CommandType.SUB_COMMANDS.value]
self.completer = AwsCompleter(
awscli_completer,
self.all_commands,
self.config,
self.config_obj,
self.log_exception,
fuzzy_match=self.get_fuzzy_match(),
shortcut_match=self.get_shortcut_match())
if refresh_resources:
self.completer.refresh_resources_and_options()
self._create_cli()
def log_exception(self, e, traceback, echo=False):
"""Logs the exception and traceback to the log file ~/.saws.log.
Args:
* e: A Exception that specifies the exception.
* traceback: A Traceback that specifies the traceback.
* echo: A boolean that specifies whether to echo the exception
to the console using click.
Returns:
None.
"""
self.logger.debug('exception: %r.', str(e))
self.logger.error("traceback: %r", traceback.format_exc())
if echo:
click.secho(str(e), fg='red')
def set_color(self, color):
"""Setter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the color flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.COLOR] = color
def get_color(self):
"""Getter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the color flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.COLOR)
def set_fuzzy_match(self, fuzzy):
"""Setter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the fuzzy flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.FUZZY] = fuzzy
self.completer.fuzzy_match = fuzzy
def get_fuzzy_match(self):
"""Getter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the fuzzy flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.FUZZY)
def set_shortcut_match(self, shortcut):
"""Setter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the shortcut flag.
Returns:
None.
"""
self.config_obj[self.config.MAIN][self.config.SHORTCUT] = shortcut
self.completer.shortcut_match = shortcut
def get_shortcut_match(self):
"""Getter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the shortcut flag.
"""
return self.config_obj[self.config.MAIN].as_bool(self.config.SHORTCUT)
def refresh_resources_and_options(self):
"""Convenience function to refresh resources and options for completion.
Used by prompt_toolkit's KeyBindingManager.
Args:
* None.
Returns:
None.
"""
self.completer.refresh_resources_and_options(force_refresh=True)
def handle_docs(self, text=None, from_fkey=False):
"""Displays contextual web docs for `F9` or the `docs` command.
Displays the web docs specific to the currently entered:
* (optional) command
* (optional) subcommand
If no command or subcommand is present, the docs index page is shown.
Docs are only displayed if:
* from_fkey is True
* from_fkey is False and `docs` is found in text
Args:
* text: A string representing the input command text.
* from_fkey: A boolean representing whether this function is
being executed from an `F9` key press.
Returns:
A boolean representing whether the web docs were shown.
"""
base_url = 'http://docs.aws.amazon.com/cli/latest/reference/'
index_html = 'index.html'
if text is None:
text = self.aws_cli.current_buffer.document.text
# If the user hit the F9 key, append 'docs' to the text
if from_fkey:
text = text.strip() + ' ' + AwsCommands.AWS_DOCS
tokens = text.split()
if len(tokens) > 2 and tokens[-1] == AwsCommands.AWS_DOCS:
prev_word = tokens[-2]
# If we have a command, build the url
if prev_word in self.commands:
prev_word = prev_word + '/'
url = base_url + prev_word + index_html
webbrowser.open(url)
return True
# if we have a command and subcommand, build the url
elif prev_word in self.sub_commands:
command_url = tokens[-3] + '/'
sub_command_url = tokens[-2] + '.html'
url = base_url + command_url + sub_command_url
webbrowser.open(url)
return True
webbrowser.open(base_url + index_html)
# If we still haven't opened the help doc at this point and the
# user hit the F9 key or typed docs, just open the main docs index
if from_fkey or AwsCommands.AWS_DOCS in tokens:
webbrowser.open(base_url + index_html)
return True
return False
def _handle_cd(self, text):
"""Handles a `cd` shell command by calling python's os.chdir.
Simply passing in the `cd` command to subprocess.call doesn't work.
Note: Changing the directory within Saws will only be in effect while
running Saws. Exiting the program will return you to the directory
you were in prior to running Saws.
Attributes:
* text: A string representing the input command text.
Returns:
A boolean representing a `cd` command was found and handled.
"""
CD_CMD = 'cd'
stripped_text = text.strip()
if stripped_text.startswith(CD_CMD):
directory = ''
if stripped_text == CD_CMD:
# Treat `cd` as a change to the root directory.
# os.path.expanduser does this in a cross platform manner.
directory = os.path.expanduser('~')
else:
tokens = text.split(CD_CMD + ' ')
directory = tokens[-1]
try:
os.chdir(directory)
except OSError as e:
self.log_exception(e, traceback, echo=True)
return True
return False
def _colorize_output(self, text):
"""Highlights output with pygments.
Only highlights the output if all of the following conditions are True:
* The color option is enabled
* The command will be handled by the `aws-cli`
* The text does not contain the `configure` command
* The text does not contain the `help` command, which already does
output highlighting
Args:
* text: A string that represents the input command text.
Returns:
A string that represents:
* The original command text if no highlighting was performed.
* The pygments highlighted command text otherwise.
"""
stripped_text = text.strip()
if not self.get_color() or stripped_text == '':
return text
if AwsCommands.AWS_COMMAND not in stripped_text.split():
return text
excludes = [AwsCommands.AWS_CONFIGURE,
AwsCommands.AWS_HELP,
'|']
if not any(substring in stripped_text for substring in excludes):
return text.strip() + self.PYGMENTS_CMD
else:
return text
def _handle_keyboard_interrupt(self, e, platform):
"""Handles keyboard interrupts more gracefully on Mac/Unix/Linux.
Allows Mac/Unix/Linux to continue running on keyboard interrupt,
as the user might interrupt a long-running AWS command with Control-C
while continuing to work with Saws.
On Windows, the "Terminate batch job (Y/N)" confirmation makes it
tricky to handle this gracefully. Thus, we re-raise KeyboardInterrupt.
Args:
* e: A KeyboardInterrupt.
* platform: A string that denotes platform such as
'Windows', 'Darwin', etc.
Returns:
None
Raises:
Exception: A KeyboardInterrupt if running on Windows.
"""
if platform == 'Windows':
raise e
else:
# Clear the renderer and send a carriage return
self.aws_cli.renderer.clear()
self.aws_cli.input_processor.feed(KeyPress(Keys.ControlM, u''))
self.aws_cli.input_processor.process_keys()
def _process_command(self, text):
"""Processes the input command, called by the cli event loop
Args:
* text: A string that represents the input command text.
Returns:
None.
"""
if AwsCommands.AWS_COMMAND in text:
text = self.completer.replace_shortcut(text)
if self.handle_docs(text):
return
try:
if not self._handle_cd(text):
text = self._colorize_output(text)
# Pass the command onto the shell so aws-cli can execute it
subprocess.call(text, shell=True)
print('')
except KeyboardInterrupt as e:
self._handle_keyboard_interrupt(e, platform.system())
except Exception as e:
self.log_exception(e, traceback, echo=True)
def _create_cli(self):
"""Creates the prompt_toolkit's CommandLineInterface.
Args:
* None.
Returns:
None.
"""
history = FileHistory(os.path.expanduser('~/.saws-history'))
toolbar = Toolbar(self.get_color,
self.get_fuzzy_match,
self.get_shortcut_match)
layout = create_default_layout(
message='saws> ',
reserve_space_for_menu=8,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar.handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
self.key_manager = KeyManager(
self.set_color,
self.get_color,
self.set_fuzzy_match,
self.get_fuzzy_match,
self.set_shortcut_match,
self.get_shortcut_match,
self.refresh_resources_and_options,
self.handle_docs)
style_factory = StyleFactory(self.theme)
application = Application(
mouse_support=False,
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.aws_cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def run_cli(self):
"""Runs the main loop.
Args:
* None.
Returns:
None.
"""
print('Version:', __version__)
print('Theme:', self.theme)
while True:
document = self.aws_cli.run(reset_current_buffer=True)
self._process_command(document.text)
def loop(cmd, history_file):
def session_toolbar(cli):
return _get_toolbar_tokens(cmd.is_conn_available, cmd.username,
cmd.connection.client.active_servers)
key_binding_manager = KeyBindingManager(
enable_search=True,
enable_abort_and_exit_bindings=True,
enable_system_bindings=True,
enable_open_in_editor=True)
bind_keys(key_binding_manager.registry)
layout = create_layout(
message='cr> ',
multiline=True,
lexer=SqlLexer,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
],
get_bottom_toolbar_tokens=session_toolbar)
application = Application(
layout=layout,
buffer=create_buffer(cmd, history_file),
style=PygmentsStyle.from_defaults(pygments_style_cls=CrateStyle),
key_bindings_registry=key_binding_manager.registry,
editing_mode=_get_editing_mode(),
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
)
eventloop = create_eventloop()
output = create_output()
cli = CommandLineInterface(application=application,
eventloop=eventloop,
output=output)
def get_num_columns_override():
return output.get_size().columns
cmd.get_num_columns = get_num_columns_override
while True:
try:
doc = cli.run(reset_current_buffer=True)
if doc:
cmd.process(doc.text)
except ProgrammingError as e:
if '401' in e.message:
username = cmd.username
password = cmd.password
cmd.username = input('Username: '******'Bye!')
return
class Haxor(object):
"""Encapsulate the Hacker News CLI.
:type cli: :class:`prompt_toolkit.CommandLineInterface`
:param cli: An instance of `prompt_toolkit.CommandLineInterface`.
:type CMDS_ENABLE_PAGINATE: list (const)
:param CMDS_ENABLE_PAGINATE: A list of commands that kick off pagination.
:type CMDS_NO_PAGINATE: list (const)
:param CMDS_NO_PAGINATE: A list of commands that disable pagination.
:type completer: :class:`prompt_toolkit.completer`
:param completer: An instance of `prompt_toolkit.completer`.
:type hacker_news_cli: :class:`hacker_news_cli.HackerNewsCli`
:param hacker_news_cli: An instance of `hacker_news_cli.HackerNewsCli`.
:type key_manager: :class:`prompt_toolkit.key_binding.manager.
KeyBindingManager`
:param key_manager: An instance of `prompt_toolkit.key_binding.manager.
KeyBindingManager`.
:type PAGINATE_CMD: str (const)
:param PAGINATE_CMD: The command to enable pagination.
:type paginate_comments: bool
:param paginate_comments: Determines whether to paginate
comments.
:type text_utils: :class:`util.TextUtils`
:param text_utils: An instance of `util.TextUtils`.
:type theme: str
:param theme: The prompt_toolkit lexer theme.
"""
CMDS_NO_PAGINATE = [
'-b',
'--browser',
'>',
'<',
]
CMDS_ENABLE_PAGINATE = [
'-cq',
'--comments_regex_query',
'-c',
'--comments',
'-cr',
'--comments_recent',
'-cu',
'--comments_unseen',
'-ch',
'--comments_hide_non_matching',
'hiring',
'freelance',
]
PAGINATE_CMD = ' | less -r'
PAGINATE_CMD_WIN = ' | more'
def __init__(self):
self.cli = None
self.key_manager = None
self.theme = 'vim'
self.paginate_comments = True
self.hacker_news_cli = HackerNewsCli()
self.text_utils = TextUtils()
self.completer = Completer(fuzzy_match=False,
text_utils=self.text_utils)
self._create_cli()
if platform.system() == 'Windows':
self.CMDS_ENABLE_PAGINATE.append('view')
def _create_key_manager(self):
"""Create the :class:`KeyManager`.
The inputs to KeyManager are expected to be callable, so we can't
use the standard @property and @attrib.setter for these attributes.
Lambdas cannot contain assignments so we're forced to define setters.
:rtype: :class:`prompt_toolkit.key_binding.manager
:return: KeyBindingManager with callables to set the toolbar options.
"""
def set_paginate_comments(paginate_comments):
"""Setter for paginating comments mode.
:type paginate: bool
:param paginate: The paginate comments mode.
"""
self.paginate_comments = paginate_comments
return KeyManager(
set_paginate_comments, lambda: self.paginate_comments)
def _create_cli(self):
"""Create the prompt_toolkit's CommandLineInterface."""
history = FileHistory(os.path.expanduser('~/.haxornewshistory'))
toolbar = Toolbar(lambda: self.paginate_comments)
layout = create_default_layout(
message=u'haxor> ',
reserve_space_for_menu=True,
get_bottom_toolbar_tokens=toolbar.handler,
)
cli_buffer = Buffer(
history=history,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
self.key_manager = self._create_key_manager()
style_factory = StyleFactory(self.theme)
application = Application(
mouse_support=False,
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def _add_comment_pagination(self, document_text):
"""Add the command to enable comment pagination where applicable.
Pagination is enabled if the command views comments and the
browser flag is not enabled.
:type document_text: str
:param document_text: The input command.
:rtype: str
:return: the input command with pagination enabled.
"""
if not any(sub in document_text for sub in self.CMDS_NO_PAGINATE):
if any(sub in document_text for sub in self.CMDS_ENABLE_PAGINATE):
if platform.system() == 'Windows':
document_text += self.PAGINATE_CMD_WIN
else:
document_text += self.PAGINATE_CMD
return document_text
def handle_exit(self, document):
"""Exits if the user typed exit or quit
:type document: :class:`prompt_toolkit.document.Document`
:param document: An instance of `prompt_toolkit.document.Document`.
"""
if document.text in ('exit', 'quit'):
sys.exit()
def run_command(self, document):
"""Run the given command.
:type document: :class:`prompt_toolkit.document.Document`
:param document: An instance of `prompt_toolkit.document.Document`.
"""
try:
if self.paginate_comments:
document.text = self._add_comment_pagination(document.text)
subprocess.call(document.text, shell=True)
except Exception as e:
click.secho(e, fg='red')
def run_cli(self):
"""Run the main loop."""
click.echo('Version: ' + __version__)
click.echo('Syntax: hn <command> [params] [options]')
while True:
document = self.cli.run()
self.handle_exit(document)
self.run_command(document)
class Haxor(object):
"""Encapsulate the Hacker News CLI.
:type cli: :class:`prompt_toolkit.CommandLineInterface`
:param cli: An instance of `prompt_toolkit.CommandLineInterface`.
:type CMDS_ENABLE_PAGINATE: list (const)
:param CMDS_ENABLE_PAGINATE: A list of commands that kick off pagination.
:type CMDS_NO_PAGINATE: list (const)
:param CMDS_NO_PAGINATE: A list of commands that disable pagination.
:type completer: :class:`prompt_toolkit.completer`
:param completer: An instance of `prompt_toolkit.completer`.
:type hacker_news_cli: :class:`hacker_news_cli.HackerNewsCli`
:param hacker_news_cli: An instance of `hacker_news_cli.HackerNewsCli`.
:type key_manager: :class:`prompt_toolkit.key_binding.manager.
KeyBindingManager`
:param key_manager: An instance of `prompt_toolkit.key_binding.manager.
KeyBindingManager`.
:type PAGINATE_CMD: str (const)
:param PAGINATE_CMD: The command to enable pagination.
:type paginate_comments: bool
:param paginate_comments: Determines whether to paginate
comments.
:type text_utils: :class:`util.TextUtils`
:param text_utils: An instance of `util.TextUtils`.
:type theme: str
:param theme: The prompt_toolkit lexer theme.
"""
CMDS_NO_PAGINATE = [
'-b',
'--browser',
'>',
'<',
]
CMDS_ENABLE_PAGINATE = [
'-cq',
'--comments_regex_query',
'-c',
'--comments',
'-cr',
'--comments_recent',
'-cu',
'--comments_unseen',
'-ch',
'--comments_hide_non_matching',
'hiring',
'freelance',
]
PAGINATE_CMD = ' | less -r'
PAGINATE_CMD_WIN = ' | more'
def __init__(self):
self.cli = None
self.key_manager = None
self.theme = 'vim'
self.paginate_comments = True
self.hacker_news_cli = HackerNewsCli()
self.text_utils = TextUtils()
self.completer = Completer(fuzzy_match=False,
text_utils=self.text_utils)
self._create_cli()
if platform.system() == 'Windows':
self.CMDS_ENABLE_PAGINATE.append('view')
def _create_key_manager(self):
"""Create the :class:`KeyManager`.
The inputs to KeyManager are expected to be callable, so we can't
use the standard @property and @attrib.setter for these attributes.
Lambdas cannot contain assignments so we're forced to define setters.
:rtype: :class:`prompt_toolkit.key_binding.manager
:return: KeyBindingManager with callables to set the toolbar options.
"""
def set_paginate_comments(paginate_comments):
"""Setter for paginating comments mode.
:type paginate: bool
:param paginate: The paginate comments mode.
"""
self.paginate_comments = paginate_comments
return KeyManager(
set_paginate_comments, lambda: self.paginate_comments)
def _create_cli(self):
"""Create the prompt_toolkit's CommandLineInterface."""
history = FileHistory(os.path.expanduser('~/.haxornewshistory'))
toolbar = Toolbar(lambda: self.paginate_comments)
layout = create_default_layout(
message=u'haxor> ',
reserve_space_for_menu=8,
get_bottom_toolbar_tokens=toolbar.handler,
)
cli_buffer = Buffer(
history=history,
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
self.key_manager = self._create_key_manager()
style_factory = StyleFactory(self.theme)
application = Application(
mouse_support=False,
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def _add_comment_pagination(self, document_text):
"""Add the command to enable comment pagination where applicable.
Pagination is enabled if the command views comments and the
browser flag is not enabled.
:type document_text: str
:param document_text: The input command.
:rtype: str
:return: the input command with pagination enabled.
"""
if not any(sub in document_text for sub in self.CMDS_NO_PAGINATE):
if any(sub in document_text for sub in self.CMDS_ENABLE_PAGINATE):
if platform.system() == 'Windows':
document_text += self.PAGINATE_CMD_WIN
else:
document_text += self.PAGINATE_CMD
return document_text
def handle_exit(self, document):
"""Exits if the user typed exit or quit
:type document: :class:`prompt_toolkit.document.Document`
:param document: An instance of `prompt_toolkit.document.Document`.
"""
if document.text in ('exit', 'quit'):
sys.exit()
def run_command(self, document):
"""Run the given command.
:type document: :class:`prompt_toolkit.document.Document`
:param document: An instance of `prompt_toolkit.document.Document`.
"""
try:
if self.paginate_comments:
text = document.text
text = self._add_comment_pagination(text)
subprocess.call(text, shell=True)
except Exception as e:
click.secho(e, fg='red')
def run_cli(self):
"""Run the main loop."""
click.echo('Version: ' + __version__)
click.echo('Syntax: hn <command> [params] [options]')
while True:
document = self.cli.run(reset_current_buffer=True)
self.handle_exit(document)
self.run_command(document)
def run(self, query=None, data=None):
self.load_config()
if data is not None or query is not None:
self.format = self.format_stdin
self.echo.verbose = False
if self.echo.verbose:
show_version()
if not self.connect():
return
if data is not None and query is None:
# cat stuff.sql | clickhouse-cli
return self.handle_input('\n'.join(data), verbose=False)
if data is None and query is not None:
# clickhouse-cli -q 'SELECT 1'
return self.handle_query(query, stream=True)
if data is not None and query is not None:
# cat stuff.csv | clickhouse-cli -q 'INSERT INTO stuff'
return self.handle_query(query, data=data, stream=True)
layout = create_prompt_layout(
lexer=PygmentsLexer(CHLexer),
get_prompt_tokens=get_prompt_tokens,
get_continuation_tokens=get_continuation_tokens,
multiline=self.multiline,
)
buffer = CLIBuffer(
client=self.client,
multiline=self.multiline,
)
application = Application(
layout=layout,
buffer=buffer,
style=CHStyle,
key_bindings_registry=KeyBinder.registry,
)
eventloop = create_eventloop()
cli = CommandLineInterface(application=application, eventloop=eventloop)
try:
while True:
try:
cli_input = cli.run(reset_current_buffer=True)
self.handle_input(cli_input.text)
except KeyboardInterrupt:
# Attempt to terminate queries
for query_id in self.query_ids:
self.client.kill_query(query_id)
self.echo.error("\nQuery was terminated.")
finally:
self.query_ids = []
except EOFError:
self.echo.success("Bye.")
class Saws(object):
"""Encapsulates the Saws CLI.
Attributes:
* aws_cli: An instance of prompt_toolkit's CommandLineInterface.
* config: An instance of Config.
* config_obj: An instance of ConfigObj, reads from ~/.sawsrc.
* aws_commands: An instance of AwsCommands
* commands: A list of commands from data/SOURCES.txt.
* sub_commands: A list of sub_commands from data/SOURCES.txt.
* global_options: A list of global_options from data/SOURCES.txt.
* resource_options: A list of resource_options from data/SOURCES.txt,
used for syntax coloring.
* ec2_states: A list of ec2_states from data/SOURCES.txt.
* completer: An instance of AwsCompleter.
* key_manager: An instance of KeyManager
* logger: An instance of SawsLogger.
* theme: A string representing the lexer theme.
Currently only 'vim' is supported.
"""
def __init__(self):
"""Inits Saws.
Args:
* None.
Returns:
None.
"""
self.aws_cli = None
self.key_manager = None
self.PYGMENTS_CMD = ' | pygmentize -l json'
self.config = Config()
self.config_obj = self.config.read_configuration()
self.theme = self.config_obj['main']['theme']
self.logger = SawsLogger(__name__,
self.config_obj['main']['log_file'],
self.config_obj['main']['log_level']).logger
self.aws_commands = AwsCommands()
self.commands, self.sub_commands, self.global_options, \
self.resource_options, self.ec2_states \
= self.aws_commands.generate_all_commands()
self.completer = AwsCompleter(
awscli_completer,
self.commands,
self.config_obj,
self.log_exception,
ec2_states=self.ec2_states,
fuzzy_match=self.get_fuzzy_match(),
shortcut_match=self.get_shortcut_match())
self.create_cli()
def log_exception(self, e, traceback, echo=False):
"""Logs the exception and traceback to the log file ~/.saws.log.
Args:
* e: A Exception that specifies the exception.
* traceback: A Traceback that specifies the traceback.
* echo: A boolean that specifies whether to echo the exception
to the console using click.
Returns:
None.
"""
self.logger.debug('exception: %r.', str(e))
self.logger.error("traceback: %r", traceback.format_exc())
if echo:
click.secho(str(e), fg='red')
def set_color(self, color):
"""Setter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the color flag.
Returns:
None.
"""
self.config_obj['main']['color_output'] = color
def get_color(self):
"""Getter for color output mode.
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the color flag.
"""
return self.config_obj['main'].as_bool('color_output')
def set_fuzzy_match(self, fuzzy):
"""Setter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the fuzzy flag.
Returns:
None.
"""
self.config_obj['main']['fuzzy_match'] = fuzzy
self.completer.fuzzy_match = fuzzy
def get_fuzzy_match(self):
"""Getter for fuzzy matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the fuzzy flag.
"""
return self.config_obj['main'].as_bool('fuzzy_match')
def set_shortcut_match(self, shortcut):
"""Setter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* color: A boolean that represents the shortcut flag.
Returns:
None.
"""
self.config_obj['main']['shortcut_match'] = shortcut
self.completer.shortcut_match = shortcut
def get_shortcut_match(self):
"""Getter for shortcut matching mode
Used by prompt_toolkit's KeyBindingManager.
KeyBindingManager expects this function to be callable so we can't use
@property and @attrib.setter.
Args:
* None.
Returns:
A boolean that represents the shortcut flag.
"""
return self.config_obj['main'].as_bool('shortcut_match')
def refresh_resources(self):
"""Convenience function to refresh resources for completion.
Used by prompt_toolkit's KeyBindingManager.
Args:
* None.
Returns:
None.
"""
self.completer.refresh_resources(force_refresh=True)
def handle_docs(self, text=None, from_fkey=False):
"""Displays contextual web docs for `F9` or the `docs` command.
Displays the web docs specific to the currently entered:
* (optional) command
* (optional) subcommand
If no command or subcommand is present, the docs index page is shown.
Docs are only displayed if:
* from_fkey is True
* from_fkey is False and `docs` is found in text
Args:
* text: A string representing the input command text.
* from_fkey: A boolean representing whether this function is
being executed from an `F9` key press.
Returns:
A boolean representing whether the web docs were shown.
"""
base_url = 'http://docs.aws.amazon.com/cli/latest/reference/'
index_html = 'index.html'
if text is None:
text = self.aws_cli.current_buffer.document.text
# If the user hit the F9 key, append 'docs' to the text
if from_fkey:
text = text.strip() + ' ' + AwsCommands.AWS_DOCS
tokens = text.split()
if len(tokens) > 2 and tokens[-1] == AwsCommands.AWS_DOCS:
prev_word = tokens[-2]
# If we have a command, build the url
if prev_word in self.commands:
prev_word = prev_word + '/'
url = base_url + prev_word + index_html
webbrowser.open(url)
return True
# if we have a command and subcommand, build the url
elif prev_word in self.sub_commands:
command_url = tokens[-3] + '/'
sub_command_url = tokens[-2] + '.html'
url = base_url + command_url + sub_command_url
webbrowser.open(url)
return True
webbrowser.open(base_url + index_html)
# If we still haven't opened the help doc at this point and the
# user hit the F9 key or typed docs, just open the main docs index
if from_fkey or AwsCommands.AWS_DOCS in tokens:
webbrowser.open(base_url + index_html)
return True
return False
def handle_cd(self, text):
"""Handles a `cd` shell command by calling python's os.chdir.
Simply passing in the `cd` command to subprocess.call doesn't work.
Note: Changing the directory within Saws will only be in effect while
running Saws. Exiting the program will return you to the directory
you were in prior to running Saws.
Attributes:
* text: A string representing the input command text.
Returns:
A boolean representing a `cd` command was found and handled.
"""
CD_CMD = 'cd'
stripped_text = text.strip()
if stripped_text.startswith(CD_CMD):
directory = ''
if stripped_text == CD_CMD:
# Treat `cd` as a change to the root directory.
# os.path.expanduser does this in a cross platform manner.
directory = os.path.expanduser('~')
else:
tokens = text.split(CD_CMD + ' ')
directory = tokens[-1]
try:
os.chdir(directory)
except OSError as e:
self.log_exception(e, traceback, echo=True)
return True
return False
def colorize_output(self, text):
"""Highlights output with pygments.
Only highlights the output if all of the following conditions are True:
* The color option is enabled
* The text does not contain the `configure` command
* The text does not contain the `help` command, which already does
output highlighting
Args:
* text: A string that represents the input command text.
Returns:
A string that represents:
* The original command text if no highlighting was performed.
* The pygments highlighted command text otherwise.
"""
if not self.get_color():
return text
stripped_text = text.strip()
excludes = [AwsCommands.AWS_CONFIGURE,
AwsCommands.AWS_HELP,
'|']
if not any(substring in stripped_text for substring in excludes):
return text.strip() + self.PYGMENTS_CMD
else:
return text
def process_command(self, text):
"""Processes the input command, called by the cli event loop
Args:
* text: A string that represents the input command text.
Returns:
None.
"""
if AwsCommands.AWS_COMMAND in text:
text = self.completer.replace_shortcut(text)
if self.handle_docs(text):
return
text = self.colorize_output(text)
try:
if not self.handle_cd(text):
# Pass the command onto the shell so aws-cli can execute it
subprocess.call(text, shell=True)
print('')
except Exception as e:
self.log_exception(e, traceback, echo=True)
def create_cli(self):
"""Creates the prompt_toolkit's CommandLineInterface.
Long description.
Args:
* None.
Returns:
None.
"""
history = FileHistory(os.path.expanduser('~/.saws-history'))
toolbar = Toolbar(self.get_color,
self.get_fuzzy_match,
self.get_shortcut_match)
layout = create_default_layout(
message='saws> ',
reserve_space_for_menu=True,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar.handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT)
self.key_manager = KeyManager(
self.set_color,
self.get_color,
self.set_fuzzy_match,
self.get_fuzzy_match,
self.set_shortcut_match,
self.get_shortcut_match,
self.refresh_resources,
self.handle_docs)
style_factory = StyleFactory(self.theme)
application = Application(
mouse_support=False,
style=style_factory.style,
layout=layout,
buffer=cli_buffer,
key_bindings_registry=self.key_manager.manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
eventloop = create_eventloop()
self.aws_cli = CommandLineInterface(
application=application,
eventloop=eventloop)
def run_cli(self):
"""Runs the main loop.
Args:
* None.
Returns:
None.
"""
print('Version:', __version__)
print('Theme:', self.theme)
while True:
document = self.aws_cli.run()
self.process_command(document.text)
class NetWork(object):
def __init__(self):
self.fuzzy = True
self.long = True
self.cmd = NetWorkCommands()
self._create_cli()
def set_fuzzy_match(self, is_fuzzy):
self.fuzzy = is_fuzzy
self.completer.set_fuzzy_match(is_fuzzy)
def get_fuzzy_match(self):
return self.fuzzy
def set_long_options(self, is_long):
self.long = is_long
self.completer.set_long_options(is_long)
def get_long_options(self):
return self.long
def _process_command(self, text):
self.cmd.run_cmd(text)
def _create_layout(self):
self.layout = create_default_layout(
message='network > ',
lexer=CommandLexer,
get_bottom_toolbar_tokens=create_toolbar_handler(
self.get_long_options, self.get_fuzzy_match))
def _create_completer(self):
self.completer = NetworkCompleter()
def _create_buffer(self):
self._create_completer()
self.buffer = Buffer(
history=FileHistory(os.path.expanduser('./network-history')),
auto_suggest=AutoSuggestFromHistory(),
enable_history_search=True,
completer=self.completer,
complete_while_typing=Always(),
accept_action=AcceptAction.RETURN_DOCUMENT,
)
def _create_style(self):
self.style = style_factory('red')
def _create_manage(self):
self.manager = get_key_manager(self.set_long_options,
self.get_long_options,
self.set_fuzzy_match,
self.get_fuzzy_match)
def _create_cli(self):
self._create_layout()
self._create_buffer()
self._create_style()
self._create_manage()
application = Application(layout=self.layout,
buffer=self.buffer,
style=self.style,
key_bindings_registry=self.manager.registry,
mouse_support=False,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
ignore_case=True)
event_loop = create_eventloop()
self.network_cli = CommandLineInterface(application=application,
eventloop=event_loop)
def _quit_command(self, text):
return (text.strip().lower() == 'exit'
or text.strip().lower() == 'quit' or text.strip() == r'\q'
or text.strip() == ':q')
def run_cli(self):
print('Version:', __version__)
while True:
document = self.network_cli.run(reset_current_buffer=True)
if self._quit_command(document.text):
raise EOFError
try:
self._process_command(document.text)
except KeyError as ex:
click.secho(ex.message, fg='red')
except NotImplementedError as ex:
click.secho(ex.message, fg='red')
class WharfeeCli(object):
"""
The CLI implementation.
"""
dcli = None
keyword_completer = None
handler = None
saved_less_opts = None
config = None
config_template = 'wharfeerc'
config_name = '~/.wharfeerc'
def __init__(self):
"""
Initialize class members.
Should read the config here at some point.
"""
self.config = self.read_configuration()
self.theme = self.config['main']['theme']
log_file = self.config['main']['log_file']
log_level = self.config['main']['log_level']
self.logger = create_logger(__name__, log_file, log_level)
# set_completer_options refreshes all by default
self.handler = DockerClient(
self.config['main'].as_int('client_timeout'),
self.clear,
self.set_completer_options)
self.completer = DockerCompleter(
long_option_names=self.get_long_options(),
fuzzy=self.get_fuzzy_match())
self.set_completer_options()
self.saved_less_opts = self.set_less_opts()
def read_configuration(self):
"""
:return:
"""
default_config = os.path.join(
self.get_package_path(), self.config_template)
write_default_config(default_config, self.config_name)
return read_config(self.config_name, default_config)
def get_package_path(self):
"""
Find out pakage root path.
:return: string: path
"""
from wharfee import __file__ as package_root
return os.path.dirname(package_root)
def set_less_opts(self):
"""
Set the "less" options and save the old settings.
What we're setting:
-F:
--quit-if-one-screen: Quit if entire file fits on first screen.
-R:
--raw-control-chars: Output "raw" control characters.
-X:
--no-init: Don't use termcap keypad init/deinit strings.
--no-keypad: Don't use termcap init/deinit strings.
This also disables launching "less" in an alternate screen.
:return: string with old options
"""
opts = os.environ.get('LESS', '')
os.environ['LESS'] = '-RXF'
return opts
def revert_less_opts(self):
"""
Restore the previous "less" options.
"""
os.environ['LESS'] = self.saved_less_opts
def write_config_file(self):
"""
Write config file on exit.
"""
self.config.write()
def clear(self):
"""
Clear the screen.
"""
click.clear()
def set_completer_options(self, cons=True, runs=True, imgs=True):
"""
Set image and container names in Completer.
Re-read if needed after a command.
:param cons: boolean: need to refresh containers
:param runs: boolean: need to refresh running containers
:param imgs: boolean: need to refresh images
"""
if cons:
cs = self.handler.containers(all=True)
if cs and len(cs) > 0 and isinstance(cs[0], dict):
containers = [name for c in cs for name in c['Names']]
self.completer.set_containers(containers)
if runs:
cs = self.handler.containers()
if cs and len(cs) > 0 and isinstance(cs[0], dict):
running = [name for c in cs for name in c['Names']]
self.completer.set_running(running)
if imgs:
def format_tagged(tagname, img_id):
if tagname == '<none>:<none>':
return img_id[:11]
return tagname
def parse_image_name(tag, img_id):
if ':' in tag:
result = tag.split(':', 2)[0]
else:
result = tag
if result == '<none>':
result = img_id[:11]
return result
ims = self.handler.images()
if ims and len(ims) > 0 and isinstance(ims[0], dict):
images = set([])
tagged = set([])
for im in ims:
repo_tag = '{0}:{1}'.format(im['Repository'], im['Tag'])
images.add(parse_image_name(repo_tag, im['Id']))
tagged.add(format_tagged(repo_tag, im['Id']))
self.completer.set_images(images)
self.completer.set_tagged(tagged)
def set_fuzzy_match(self, is_fuzzy):
"""
Setter for fuzzy matching mode
:param is_fuzzy: boolean
"""
self.config['main']['fuzzy_match'] = is_fuzzy
self.completer.set_fuzzy_match(is_fuzzy)
def get_fuzzy_match(self):
"""
Getter for fuzzy matching mode
:return: boolean
"""
return self.config['main'].as_bool('fuzzy_match')
def set_long_options(self, is_long):
"""
Setter for long option names.
:param is_long: boolean
"""
self.config['main']['suggest_long_option_names'] = is_long
self.completer.set_long_options(is_long)
def get_long_options(self):
"""
Getter for long option names.
:return: boolean
"""
return self.config['main'].as_bool('suggest_long_option_names')
def refresh_completions(self):
"""
After processing the command, refresh the lists of
containers and images as needed
"""
self.set_completer_options(self.handler.is_refresh_containers,
self.handler.is_refresh_running,
self.handler.is_refresh_images)
def run_cli(self):
"""
Run the main loop
"""
print('Version:', __version__)
print('Home: http://wharfee.com')
history = FileHistory(os.path.expanduser('~/.wharfee-history'))
toolbar_handler = create_toolbar_handler(self.get_long_options, self.get_fuzzy_match)
layout = create_default_layout(
message='wharfee> ',
reserve_space_for_menu=True,
lexer=CommandLexer,
get_bottom_toolbar_tokens=toolbar_handler,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = Buffer(
history=history,
completer=self.completer,
complete_while_typing=Always())
manager = get_key_manager(
self.set_long_options,
self.get_long_options,
self.set_fuzzy_match,
self.get_fuzzy_match)
application = Application(
style=style_factory(self.theme),
layout=layout,
buffer=cli_buffer,
key_bindings_registry=manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
ignore_case=True)
eventloop = create_eventloop()
self.dcli = CommandLineInterface(
application=application,
eventloop=eventloop)
while True:
try:
document = self.dcli.run()
self.handler.handle_input(document.text)
if isinstance(self.handler.output, GeneratorType):
output_stream(self.handler.command,
self.handler.output,
self.handler.logs)
elif self.handler.output is not None:
lines = format_data(
self.handler.command,
self.handler.output)
click.echo_via_pager('\n'.join(lines))
if self.handler.after:
for line in self.handler.after():
click.echo(line)
self.refresh_completions()
except OptionError as ex:
self.logger.debug('Error: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(ex.msg, fg='red')
except KeyboardInterrupt:
# user pressed Ctrl + C
if self.handler.after:
click.echo('')
for line in self.handler.after():
click.echo(line)
self.refresh_completions()
except DockerPermissionException as ex:
self.logger.debug('Permission exception: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho(ex.message, fg='red')
except EOFError:
# exit out of the CLI
break
# TODO: uncomment for release
except Exception as ex:
self.logger.debug('Exception: %r.', ex)
self.logger.error("traceback: %r", traceback.format_exc())
click.secho("{0}".format(ex), fg='red')
break
self.revert_less_opts()
self.write_config_file()
print('Goodbye!')
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')
on_error_modes = {'STOP': ON_ERROR_STOP, 'RESUME': ON_ERROR_RESUME}
self.on_error = on_error_modes[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):
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, 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,
on_error=self.on_error)
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:
successful = True
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(
need_completion_reset(document.text))
# 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, 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 run_cli(self):
pgexecute = self.pgexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
completer = self.completer
self.refresh_completions()
def set_vi_mode(value):
self.vi_mode = value
key_binding_manager = pgcli_bindings(
get_vi_mode_enabled=lambda: self.vi_mode,
set_vi_mode_enabled=set_vi_mode)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/pgcli')
print('Mail: https://groups.google.com/forum/#!forum/pgcli')
print('Home: http://pgcli.com')
def prompt_tokens(cli):
return [(Token.Prompt, '%s> ' % pgexecute.dbname)]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.vi_mode)
layout = create_default_layout(lexer=PostgresLexer,
reserve_space_for_menu=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
extra_input_processors=[
# Highlight matching brackets while editing.
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
])
history_file = self.config['main']['history_file']
buf = PGBuffer(always_multiline=self.multi_line, completer=completer,
history=FileHistory(os.path.expanduser(history_file)),
complete_while_typing=Always())
application = Application(style=style_factory(self.syntax_style),
layout=layout, buffer=buf,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION)
cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = cli.run()
# The reason we check here instead of inside the pgexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the pgexecute.run()
# statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
successful = False
# Initialized to [] because res might never get initialized
# if an exception occurs in pgexecute.run(). Which causes
# finally clause to fail.
res = []
start = time()
# Run the query.
res = pgexecute.run(document.text, self.pgspecial)
duration = time() - start
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and
cur and cur.rowcount > threshold):
click.secho('The result set has more than %s rows.'
% threshold, fg='red')
if not click.confirm('Do you want to continue?'):
click.secho("Aborted!", err=True, fg='red')
break
formatted = format_output(title, cur, headers, status,
self.table_format,
self.pgspecial.expanded_output)
output.extend(formatted)
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
pgexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
click.secho("cancelled query", err=True, fg='red')
except NotImplementedError:
click.secho('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
reconnect = True
if ('server closed the connection' in utf8tounicode(e.args[0])):
reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
show_default=False, type=bool, default=True)
if reconnect:
try:
pgexecute.connect()
click.secho('Reconnected!\nTry the command again.', fg='green')
except OperationalError as e:
click.secho(str(e), err=True, fg='red')
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
click.secho(str(e), err=True, fg='red')
else:
click.echo_via_pager('\n'.join(output))
if self.pgspecial.timing_enabled:
print('Command Time: %0.03fs' % duration)
print('Format Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions()
# Refresh search_path to set default schema.
if need_search_path_refresh(document.text):
logger.debug('Refreshing search path')
completer.set_search_path(pgexecute.search_path())
logger.debug('Search path: %r', completer.search_path)
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
print ('Goodbye!')
finally: # Reset the less opts back to original.
logger.debug('Restoring env var LESS to %r.', original_less_opts)
os.environ['LESS'] = original_less_opts
def run_cli(self):
sqlexecute = self.sqlexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
self.initialize_completions()
completer = self.completer
def set_key_bindings(value):
if value not in ('emacs', 'vi'):
value = 'emacs'
self.key_bindings = value
project_root = os.path.dirname(PACKAGE_ROOT)
author_file = os.path.join(project_root, 'AUTHORS')
sponsor_file = os.path.join(project_root, 'SPONSORS')
key_binding_manager = mycli_bindings(
get_key_bindings=lambda: self.key_bindings,
set_key_bindings=set_key_bindings)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/mycli')
print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
print('Home: http://mycli.net')
print('Thanks to the contributor -',
thanks_picker([author_file, sponsor_file]))
def prompt_tokens(cli):
return [(Token.Prompt, self.get_prompt(self.prompt_format))]
get_toolbar_tokens = create_toolbar_tokens_func(
lambda: self.key_bindings)
layout = create_default_layout(
lexer=MyCliLexer,
reserve_space_for_menu=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(
chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
])
buf = CLIBuffer(always_multiline=self.multi_line,
completer=completer,
history=FileHistory(
os.path.expanduser('~/.mycli-history')),
complete_while_typing=Always())
application = Application(
style=style_factory(self.syntax_style),
layout=layout,
buffer=buf,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION)
cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = cli.run()
special.set_expanded_output(False)
# The reason we check here instead of inside the sqlexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the
# sqlexecute.run() statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
continue
destroy = confirm_destructive_query(document.text)
if destroy is None:
pass # Query was not destructive. Nothing to do here.
elif destroy is True:
self.output('Your call!')
else:
self.output('Wise choice!')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
if self.logfile:
self.logfile.write('\n# %s\n' % datetime.now())
self.logfile.write(document.text)
self.logfile.write('\n')
successful = False
start = time()
res = sqlexecute.run(document.text)
duration = time() - start
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and cur
and cur.rowcount > threshold):
self.output(
'The result set has more than %s rows.' %
threshold,
fg='red')
if not click.confirm('Do you want to continue?'):
self.output("Aborted!", err=True, fg='red')
break
output.extend(
format_output(title, cur, headers, status,
self.table_format))
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
sqlexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
self.output("cancelled query", err=True, fg='red')
except NotImplementedError:
self.output('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
logger.debug("Exception: %r", e)
reconnect = True
if (e.args[0] in (2003, 2006, 2013)):
reconnect = click.prompt(
'Connection reset. Reconnect (Y/n)',
show_default=False,
type=bool,
default=True)
if reconnect:
logger.debug('Attempting to reconnect.')
try:
sqlexecute.connect()
logger.debug('Reconnected successfully.')
self.output(
'Reconnected!\nTry the command again.',
fg='green')
except OperationalError as e:
logger.debug('Reconnect failed. e: %r', e)
self.output(str(e), err=True, fg='red')
continue # If reconnection failed, don't proceed further.
else: # If user chooses not to reconnect, don't proceed further.
continue
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
else:
self.output_via_pager('\n'.join(output))
if special.is_timing_enabled():
self.output('Command Time: %0.03fs' % duration)
self.output('Format Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_dynamic_completions()
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
self.output('Goodbye!')
finally: # Reset the less opts back to original.
logger.debug('Restoring env var LESS to %r.', original_less_opts)
os.environ['LESS'] = original_less_opts
os.environ['PAGER'] = special.get_original_pager()
class CLI:
def __init__(self, host, port, user, password, database, settings, format,
format_stdin, multiline, stacktrace):
self.config = None
self.host = host
self.port = port
self.user = user
self.password = password
self.database = database
self.settings = {k: v[0] for k, v in parse_qs(settings).items()}
self.format = format
self.format_stdin = format_stdin
self.multiline = multiline
self.stacktrace = stacktrace
self.server_version = None
self.query_ids = []
self.client = None
self.echo = Echo(verbose=True, colors=True)
self.progress = None
self.metadata = {}
def connect(self):
self.scheme = 'http'
if '://' in self.host:
u = urlparse(self.host, allow_fragments=False)
self.host = u.hostname
self.port = u.port or self.port
self.scheme = u.scheme
self.url = '{scheme}://{host}:{port}/'.format(scheme=self.scheme,
host=self.host,
port=self.port)
self.client = Client(
self.url,
self.user,
self.password,
self.database,
self.settings,
self.stacktrace,
self.conn_timeout,
self.conn_timeout_retry,
self.conn_timeout_retry_delay,
)
self.echo.print("Connecting to {host}:{port}".format(host=self.host,
port=self.port))
try:
response = self.client.query('SELECT version();',
fmt='TabSeparated')
except TimeoutError:
self.echo.error("Error: Connection timeout.")
return False
except ConnectionError:
self.echo.error("Error: Failed to connect.")
return False
except DBException as e:
self.echo.error("Error:")
self.echo.error(e.error)
if self.stacktrace and e.stacktrace:
self.echo.print("Stack trace:")
self.echo.print(e.stacktrace)
return False
if not response.data.endswith('\n'):
self.echo.error(
"Error: Request failed: `SELECT version();` query failed.")
return False
version = response.data.strip().split('.')
self.server_version = (int(version[0]), int(version[1]),
int(version[2]))
self.echo.success(
"Connected to ClickHouse server v{0}.{1}.{2}.\n".format(
*self.server_version))
return True
def load_config(self):
self.config = read_config()
self.multiline = self.config.getboolean('main', 'multiline')
self.format = self.format or self.config.get('main', 'format')
self.format_stdin = self.format_stdin or self.config.get(
'main', 'format_stdin')
self.show_formatted_query = self.config.getboolean(
'main', 'show_formatted_query')
self.highlight = self.config.getboolean('main', 'highlight')
self.highlight_output = self.config.getboolean('main',
'highlight_output')
self.highlight_truecolor = self.config.getboolean(
'main', 'highlight_truecolor') and os.environ.get('COLORTERM')
self.conn_timeout = self.config.getfloat('http', 'conn_timeout')
self.conn_timeout_retry = self.config.getint('http',
'conn_timeout_retry')
self.conn_timeout_retry_delay = self.config.getfloat(
'http', 'conn_timeout_retry_delay')
self.host = self.host or self.config.get('defaults',
'host') or '127.0.0.1'
self.port = self.port or self.config.get('defaults', 'port') or 8123
self.user = self.user or self.config.get('defaults',
'user') or 'default'
self.password = self.password or self.config.get(
'defaults', 'password') or ''
self.database = self.database or self.config.get('defaults',
'db') or 'default'
config_settings = dict(self.config.items('settings'))
arg_settings = self.settings
config_settings.update(arg_settings)
self.settings = config_settings
self.echo.colors = self.highlight
def run(self, query, data):
self.load_config()
if data or query is not None:
self.format = self.format_stdin
self.echo.verbose = False
if self.echo.verbose:
show_version()
if not self.connect():
return
if self.client:
self.client.settings = self.settings
self.client.cli_settings = {
'multiline': self.multiline,
'format': self.format,
'format_stdin': self.format_stdin,
'show_formatted_query': self.show_formatted_query,
'highlight': self.highlight,
'highlight_output': self.highlight_output,
}
if data and query is None:
# cat stuff.sql | clickhouse-cli
# clickhouse-cli stuff.sql
for subdata in data:
self.handle_input(subdata.read(),
verbose=False,
refresh_metadata=False)
return
if not data and query is not None:
# clickhouse-cli -q 'SELECT 1'
return self.handle_query(query, stream=True)
if data and query is not None:
# cat stuff.csv | clickhouse-cli -q 'INSERT INTO stuff'
# clickhouse-cli -q 'INSERT INTO stuff' stuff.csv
for subdata in data:
compress = 'gzip' if os.path.splitext(
subdata.name)[1] == '.gz' else False
self.handle_query(query,
data=subdata,
stream=True,
compress=compress)
return
layout = create_prompt_layout(
lexer=PygmentsLexer(CHLexer) if self.highlight else None,
get_prompt_tokens=get_prompt_tokens,
get_continuation_tokens=get_continuation_tokens,
multiline=self.multiline,
)
buffer = CLIBuffer(
client=self.client,
multiline=self.multiline,
metadata=self.metadata,
)
application = Application(
layout=layout,
buffer=buffer,
style=CHStyle if self.highlight else None,
key_bindings_registry=KeyBinder.registry,
)
eventloop = create_eventloop()
self.cli = CommandLineInterface(application=application,
eventloop=eventloop)
self.cli.application.buffer.completer.refresh_metadata()
try:
while True:
try:
cli_input = self.cli.run(reset_current_buffer=True)
self.handle_input(cli_input.text)
except KeyboardInterrupt:
# Attempt to terminate queries
for query_id in self.query_ids:
self.client.kill_query(query_id)
self.echo.error("\nQuery was terminated.")
finally:
self.query_ids = []
except EOFError:
self.echo.success("Bye.")
def handle_input(self, input_data, verbose=True, refresh_metadata=True):
force_pager = False
if input_data.endswith(
r'\p' if isinstance(input_data, str) else rb'\p'):
input_data = input_data[:-2]
force_pager = True
# FIXME: A dirty dirty hack to make multiple queries (per one paste) work.
self.query_ids = []
for query in sqlparse.split(input_data):
query_id = str(uuid4())
self.query_ids.append(query_id)
self.handle_query(query,
verbose=verbose,
query_id=query_id,
force_pager=force_pager)
if refresh_metadata and input_data:
self.cli.application.buffer.completer.refresh_metadata()
def handle_query(self,
query,
data=None,
stream=False,
verbose=False,
query_id=None,
compress=False,
**kwargs):
if query.rstrip(';') == '':
return
elif query.lower() in EXIT_COMMANDS:
raise EOFError
elif query.lower() in (r'\?', 'help'):
rows = [
['', ''],
["clickhouse-cli's custom commands:", ''],
['---------------------------------', ''],
['USE', "Change the current database."],
['SET', "Set an option for the current CLI session."],
['QUIT', "Exit clickhouse-cli."],
['HELP', "Show this help message."],
['', ''],
["PostgreSQL-like custom commands:", ''],
['--------------------------------', ''],
[r'\l', "Show databases."],
[r'\c', "Change the current database."],
[r'\d, \dt', "Show tables in the current database."],
[r'\d+', "Show table's schema."],
[r'\ps', "Show current queries."],
[r'\kill', "Kill query by its ID."],
['', ''],
["Query suffixes:", ''],
['---------------', ''],
[r'\g, \G', "Use the Vertical format."],
[r'\p', "Enable the pager."],
]
for row in rows:
self.echo.success('{:<8s}'.format(row[0]), nl=False)
self.echo.info(row[1])
return
elif query in (r'\d', r'\dt'):
query = 'SHOW TABLES'
elif query.startswith(r'\d+ '):
query = 'DESCRIBE TABLE ' + query[4:]
elif query == r'\l':
query = 'SHOW DATABASES'
elif query.startswith(r'\c '):
query = 'USE ' + query[3:]
elif query.startswith(r'\ps'):
query = (
"SELECT query_id, user, address, elapsed, {}, memory_usage "
"FROM system.processes WHERE query_id != '{}'").format(
'read_rows' if self.server_version[2] >= 54115 else
'rows_read', query_id)
elif query.startswith(r'\kill '):
self.client.kill_query(query[6:])
return
response = ''
self.progress_reset()
try:
response = self.client.query(
query,
fmt=self.format,
data=data,
stream=stream,
verbose=verbose,
query_id=query_id,
compress=compress,
)
except TimeoutError:
self.echo.error("Error: Connection timeout.")
return
except ConnectionError:
self.echo.error("Error: Failed to connect.")
return
except DBException as e:
self.progress_reset()
self.echo.error("\nQuery:")
self.echo.error(query)
self.echo.error("\n\nReceived exception from server:")
self.echo.error(e.error)
if self.stacktrace and e.stacktrace:
self.echo.print("\nStack trace:")
self.echo.print(e.stacktrace)
self.echo.print('\nElapsed: {elapsed:.3f} sec.\n'.format(
elapsed=e.response.elapsed.total_seconds()))
return
total_rows, total_bytes = self.progress_reset()
self.echo.print()
if stream:
data = response.iter_lines() if hasattr(
response, 'iter_lines') else response.data
for line in data:
print(line.decode('utf-8', 'ignore'))
else:
if response.data != '':
print_func = print
if self.config.getboolean('main', 'pager') or kwargs.pop(
'force_pager', False):
print_func = self.echo.pager
should_highlight_output = (verbose and self.highlight
and self.highlight_output and
response.format in PRETTY_FORMATS)
formatter = TerminalFormatter()
if self.highlight and self.highlight_truecolor:
formatter = TerminalTrueColorFormatter(
style=CHPygmentsStyle)
if should_highlight_output:
print_func(
pygments.highlight(response.data,
CHPrettyFormatLexer(), formatter))
else:
print_func(response.data)
if response.message != '':
self.echo.print(response.message)
self.echo.print()
self.echo.success('Ok. ', nl=False)
if response.rows is not None:
self.echo.print('{rows_count} row{rows_plural} in set.'.format(
rows_count=response.rows,
rows_plural='s' if response.rows != 1 else '',
),
end=' ')
if self.config.getboolean(
'main', 'timing') and response.time_elapsed is not None:
self.echo.print(
'Elapsed: {elapsed:.3f} sec. Processed: {rows} rows, {bytes} ({avg_rps} rows/s, {avg_bps}/s)'
.format(
elapsed=response.time_elapsed,
rows=numberunit_fmt(total_rows),
bytes=sizeof_fmt(total_bytes),
avg_rps=numberunit_fmt(total_rows /
max(response.time_elapsed, 0.001)),
avg_bps=sizeof_fmt(total_bytes /
max(response.time_elapsed, 0.001)),
),
end='')
self.echo.print('\n')
def progress_update(self, line):
if not self.config.getboolean('main', 'timing'):
return
# Parse X-ClickHouse-Progress header
now = datetime.now()
progress = json.loads(line[23:].decode().strip())
progress = {
'timestamp': now,
'read_rows': int(progress['read_rows']),
'total_rows': int(progress['total_rows']),
'read_bytes': int(progress['read_bytes']),
}
# Calculate percentage completed and format initial message
progress['percents'] = int(
(progress['read_rows'] / progress['total_rows']) *
100) if progress['total_rows'] > 0 else 0
message = 'Progress: {} rows, {}'.format(
numberunit_fmt(progress['read_rows']),
sizeof_fmt(progress['read_bytes']))
# Calculate row and byte read velocity
if self.progress:
delta = (now - self.progress['timestamp']).total_seconds()
if delta > 0:
rps = (progress['read_rows'] -
self.progress['read_rows']) / delta
bps = (progress['read_bytes'] -
self.progress['read_bytes']) / delta
message += ' ({} rows/s, {}/s)'.format(numberunit_fmt(rps),
sizeof_fmt(bps))
self.progress = progress
self.progress_print(message, progress['percents'])
def progress_reset(self):
progress = self.progress
self.progress = None
clickhouse_cli.helpers.trace_headers_stream = self.progress_update
# Clear printed progress (if any)
columns = shutil.get_terminal_size((80, 0)).columns
sys.stdout.write(u"\u001b[%dD" % columns + " " * columns)
sys.stdout.flush()
# Report totals
if progress:
return (progress['read_rows'], progress['read_bytes'])
return (0, 0)
def progress_print(self, message, percents):
suffix = '%3d%%' % percents
columns = shutil.get_terminal_size((80, 0)).columns
bars_max = columns - (len(message) + len(suffix) + 3)
bars = int(percents * (bars_max / 100)) if (bars_max > 0) else 0
message = '{} \033[42m{}\033[0m{} {}'.format(message, " " * bars,
" " * (bars_max - bars),
suffix)
sys.stdout.write(u"\u001b[%dD" % columns + message)
sys.stdout.flush()
def loop(cmd, history_file):
key_binding_manager = KeyBindingManager(
enable_search=True,
enable_abort_and_exit_bindings=True,
enable_system_bindings=True,
enable_open_in_editor=True
)
bind_keys(key_binding_manager.registry)
layout = create_layout(
multiline=True,
lexer=SqlLexer,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
],
get_bottom_toolbar_tokens=lambda cli: get_toolbar_tokens(cmd),
get_prompt_tokens=lambda cli: [(Token.Prompt, 'cr> ')]
)
application = Application(
layout=layout,
buffer=create_buffer(cmd, history_file),
style=PygmentsStyle.from_defaults(pygments_style_cls=CrateStyle),
key_bindings_registry=key_binding_manager.registry,
editing_mode=_get_editing_mode(),
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
)
eventloop = create_eventloop()
output = create_output()
cli = CommandLineInterface(
application=application,
eventloop=eventloop,
output=output
)
def get_num_columns_override():
return output.get_size().columns
cmd.get_num_columns = get_num_columns_override
while True:
try:
doc = cli.run(reset_current_buffer=True)
if doc:
cmd.process(doc.text)
except ProgrammingError as e:
if '401' in e.message:
username = cmd.username
password = cmd.password
cmd.username = input('Username: '******'Bye!')
return
class MyCli(object):
default_prompt = '\\t \\u@\\h:\\d> '
defaults_suffix = None
# In order of being loaded. Files lower in list override earlier ones.
cnf_files = [
'/etc/my.cnf',
'/etc/mysql/my.cnf',
'/usr/local/etc/my.cnf',
'~/.my.cnf'
]
system_config_files = [
'/etc/myclirc',
]
default_config_file = os.path.join(PACKAGE_ROOT, 'myclirc')
user_config_file = '~/.myclirc'
def __init__(self, sqlexecute=None, prompt=None,
logfile=None, defaults_suffix=None, defaults_file=None,
login_path=None, auto_vertical_output=False):
self.sqlexecute = sqlexecute
self.logfile = logfile
self.defaults_suffix = defaults_suffix
self.login_path = login_path
self.auto_vertical_output = auto_vertical_output
# self.cnf_files is a class variable that stores the list of mysql
# config files to read in at launch.
# If defaults_file is specified then override the class variable with
# defaults_file.
if defaults_file:
self.cnf_files = [defaults_file]
# Load config.
config_files = ([self.default_config_file] + self.system_config_files +
[self.user_config_file])
c = self.config = read_config_files(config_files)
self.multi_line = c['main'].as_bool('multi_line')
self.destructive_warning = c['main'].as_bool('destructive_warning')
self.key_bindings = c['main']['key_bindings']
special.set_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')
# Write user config if system config wasn't the last config loaded.
if c.filename not in self.system_config_files:
write_default_config(self.default_config_file, self.user_config_file)
# audit log
if self.logfile is None and 'audit_log' in c['main']:
try:
self.logfile = open(os.path.expanduser(c['main']['audit_log']), 'a')
except (IOError, OSError) as e:
self.output('Error: Unable to open the audit log file. Your queries will not be logged.', err=True, fg='red')
self.logfile = False
self.completion_refresher = CompletionRefresher()
self.logger = logging.getLogger(__name__)
self.initialize_logging()
prompt_cnf = self.read_my_cnf_files(self.cnf_files, ['prompt'])['prompt']
self.prompt_format = prompt or prompt_cnf or c['main']['prompt'] or \
self.default_prompt
self.query_history = []
# Initialize completer.
smart_completion = c['main'].as_bool('smart_completion')
self.completer = SQLCompleter(smart_completion)
self._completer_lock = threading.Lock()
# Register custom special commands.
self.register_special_commands()
# Load .mylogin.cnf if it exists.
mylogin_cnf_path = get_mylogin_cnf_path()
if mylogin_cnf_path:
try:
mylogin_cnf = open_mylogin_cnf(mylogin_cnf_path)
if mylogin_cnf_path and mylogin_cnf:
# .mylogin.cnf gets read last, even if defaults_file is specified.
self.cnf_files.append(mylogin_cnf)
elif mylogin_cnf_path and not mylogin_cnf:
# There was an error reading the login path file.
print('Error: Unable to read login path file.')
except CryptoError:
click.secho('Warning: .mylogin.cnf was not read: pycrypto '
'module is not available.')
self.cli = None
def register_special_commands(self):
special.register_special_command(self.change_db, 'use',
'\\u', 'Change to a new database.', aliases=('\\u',))
special.register_special_command(self.change_db, 'connect',
'\\r', 'Reconnect to the database. Optional database argument.',
aliases=('\\r', ), case_sensitive=True)
special.register_special_command(self.refresh_completions, 'rehash',
'\\#', 'Refresh auto-completions.', arg_type=NO_QUERY, aliases=('\\#',))
special.register_special_command(self.change_table_format, 'tableformat',
'\\T', 'Change Table Type.', aliases=('\\T',), case_sensitive=True)
special.register_special_command(self.execute_from_file, 'source', '\\. filename',
'Execute commands from file.', aliases=('\\.',))
special.register_special_command(self.change_prompt_format, 'prompt',
'\\R', 'Change prompt format.', aliases=('\\R',), case_sensitive=True)
def change_table_format(self, arg, **_):
if not arg in table_formats():
msg = "Table type %s not yet implemented. Allowed types:" % arg
for table_type in table_formats():
msg += "\n\t%s" % table_type
yield (None, None, None, msg)
else:
self.table_format = arg
yield (None, None, None, "Changed table Type to %s" % self.table_format)
def change_db(self, arg, **_):
if arg is None:
self.sqlexecute.connect()
else:
self.sqlexecute.connect(database=arg)
yield (None, None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.sqlexecute.dbname, self.sqlexecute.user))
def execute_from_file(self, arg, **_):
if not arg:
message = 'Missing required argument, filename.'
return [(None, None, None, message)]
try:
with open(os.path.expanduser(arg), encoding='utf-8') as f:
query = f.read()
except IOError as e:
return [(None, None, None, str(e))]
return self.sqlexecute.run(query)
def change_prompt_format(self, arg, **_):
"""
Change the prompt format.
"""
if not arg:
message = 'Missing required argument, format.'
return [(None, None, None, message)]
self.prompt_format = self.get_prompt(arg)
return [(None, None, None, "Changed prompt format to %s" % arg)]
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('mycli')
root_logger.addHandler(handler)
root_logger.setLevel(level_map[log_level.upper()])
logging.captureWarnings(True)
root_logger.debug('Initializing mycli logging.')
root_logger.debug('Log file %r.', log_file)
def connect_uri(self, uri):
uri = urlparse(uri)
database = uri.path[1:] # ignore the leading fwd slash
self.connect(database, uri.username, uri.password, uri.hostname,
uri.port)
def read_my_cnf_files(self, files, keys):
"""
Reads a list of config files and merges them. The last one will win.
:param files: list of files to read
:param keys: list of keys to retrieve
:returns: tuple, with None for missing keys.
"""
cnf = read_config_files(files)
sections = ['client']
if self.login_path and self.login_path != 'client':
sections.append(self.login_path)
if self.defaults_suffix:
sections.extend([sect + self.defaults_suffix for sect in sections])
def get(key):
result = None
for sect in cnf:
if sect in sections and key in cnf[sect]:
result = cnf[sect][key]
return result
return dict([(x, get(x)) for x in keys])
def connect(self, database='', user='', passwd='', host='', port='',
socket='', charset=''):
cnf = {'database': None,
'user': None,
'password': None,
'host': None,
'port': None,
'socket': None,
'default-character-set': None}
cnf = self.read_my_cnf_files(self.cnf_files, cnf.keys())
# Fall back to config values only if user did not specify a value.
database = database or cnf['database']
if port or host:
socket = ''
else:
socket = socket or cnf['socket']
user = user or cnf['user'] or os.getenv('USER')
host = host or cnf['host'] or 'localhost'
port = port or cnf['port'] or 3306
try:
port = int(port)
except ValueError as e:
self.output("Error: Invalid port number: '{0}'.".format(port),
err=True, fg='red')
exit(1)
passwd = passwd or cnf['password']
charset = charset or cnf['default-character-set'] or 'utf8'
# Connect to the database.
try:
try:
sqlexecute = SQLExecute(database, user, passwd, host, port,
socket, charset)
except OperationalError as e:
if ('Access denied for user' in e.args[1]):
passwd = click.prompt('Password', hide_input=True,
show_default=False, type=str)
sqlexecute = SQLExecute(database, user, passwd, host, port,
socket, charset)
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())
self.output(str(e), err=True, fg='red')
exit(1)
self.sqlexecute = sqlexecute
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):
sqlexecute = self.sqlexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
self.set_pager_from_config()
self.refresh_completions()
def set_key_bindings(value):
if value not in ('emacs', 'vi'):
value = 'emacs'
self.key_bindings = value
project_root = os.path.dirname(PACKAGE_ROOT)
author_file = os.path.join(project_root, 'AUTHORS')
sponsor_file = os.path.join(project_root, 'SPONSORS')
key_binding_manager = mycli_bindings(get_key_bindings=lambda: self.key_bindings,
set_key_bindings=set_key_bindings)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/mycli')
print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
print('Home: http://mycli.net')
print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file]))
def prompt_tokens(cli):
return [(Token.Prompt, self.get_prompt(self.prompt_format))]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.key_bindings,
self.completion_refresher.is_refreshing)
layout = create_prompt_layout(lexer=MyCliLexer,
multiline=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
display_completions_in_columns=self.wider_completion_menu,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
])
with self._completer_lock:
buf = CLIBuffer(always_multiline=self.multi_line, completer=self.completer,
history=FileHistory(os.path.expanduser('~/.mycli-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)
self.cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = self.cli.run()
special.set_expanded_output(False)
# The reason we check here instead of inside the sqlexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the
# sqlexecute.run() statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(self.cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
continue
if self.destructive_warning:
destroy = confirm_destructive_query(document.text)
if destroy is None:
pass # Query was not destructive. Nothing to do here.
elif destroy is True:
self.output('Your call!')
else:
self.output('Wise choice!')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
if self.logfile:
self.logfile.write('\n# %s\n' % datetime.now())
self.logfile.write(document.text)
self.logfile.write('\n')
successful = False
start = time()
res = sqlexecute.run(document.text)
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
threshold = 1000
if (is_select(status) and
cur and cur.rowcount > threshold):
self.output('The result set has more than %s rows.'
% threshold, fg='red')
if not click.confirm('Do you want to continue?'):
self.output("Aborted!", err=True, fg='red')
break
if self.auto_vertical_output:
max_width = self.cli.output.get_size().columns
else:
max_width = None
formatted = format_output(title, cur, headers,
status, self.table_format,
special.is_expanded_output(), max_width)
output.extend(formatted)
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except UnicodeDecodeError as e:
import pymysql
if pymysql.VERSION < ('0', '6', '7'):
message = ('You are running an older version of pymysql.\n'
'Please upgrade to 0.6.7 or above to view binary data.\n'
'Try \'pip install -U pymysql\'.')
self.output(message)
else:
raise e
except KeyboardInterrupt:
# Restart connection to the database
sqlexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
self.output("cancelled query", err=True, fg='red')
except NotImplementedError:
self.output('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
logger.debug("Exception: %r", e)
reconnect = True
if (e.args[0] in (2003, 2006, 2013)):
reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
show_default=False, type=bool, default=True)
if reconnect:
logger.debug('Attempting to reconnect.')
try:
sqlexecute.connect()
logger.debug('Reconnected successfully.')
self.output('Reconnected!\nTry the command again.', fg='green')
except OperationalError as e:
logger.debug('Reconnect failed. e: %r', e)
self.output(str(e), err=True, fg='red')
continue # If reconnection failed, don't proceed further.
else: # If user chooses not to reconnect, don't proceed further.
continue
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
else:
try:
if special.is_pager_enabled():
self.output_via_pager('\n'.join(output))
else:
self.output('\n'.join(output))
except KeyboardInterrupt:
pass
if special.is_timing_enabled():
self.output('Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions(
reset=need_completion_reset(document.text))
finally:
if self.logfile is False:
self.output("Warning: This query was not logged.", err=True, fg='red')
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
self.output('Goodbye!')
finally: # Reset the less opts back to original.
logger.debug('Restoring env var LESS to %r.', original_less_opts)
os.environ['LESS'] = original_less_opts
os.environ['PAGER'] = special.get_original_pager()
def output(self, text, **kwargs):
if self.logfile:
self.logfile.write(utf8tounicode(text))
self.logfile.write('\n')
click.secho(text, **kwargs)
def output_via_pager(self, text):
if self.logfile:
self.logfile.write(text)
self.logfile.write('\n')
click.echo_via_pager(text)
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 set_pager_from_config(self):
cnf = self.read_my_cnf_files(self.cnf_files, ['pager', 'skip-pager'])
if cnf['pager']:
special.set_pager(cnf['pager'])
if cnf['skip-pager']:
special.disable_pager()
def refresh_completions(self, reset=False):
if reset:
with self._completer_lock:
self.completer.reset_completions()
self.completion_refresher.refresh(self.sqlexecute,
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 mycli 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):
sqlexecute = self.sqlexecute
string = string.replace('\\u', sqlexecute.user or '(none)')
string = string.replace('\\h', sqlexecute.host or '(none)')
string = string.replace('\\d', sqlexecute.dbname or '(none)')
string = string.replace('\\t', sqlexecute.server_type()[0] or 'mycli')
string = string.replace('\\n', "\n")
return string
class VCli(object):
def __init__(self, vexecute=None, vclirc_file=None):
self.vexecute = vexecute
from vcli import __file__ as package_root
package_root = os.path.dirname(package_root)
default_config = os.path.join(package_root, 'vclirc')
write_default_config(default_config, vclirc_file)
self.vspecial = VSpecial()
# Load config.
c = self.config = load_config(vclirc_file, default_config)
self.multi_line = c['main'].as_bool('multi_line')
self.vi_mode = c['main'].as_bool('vi')
self.vspecial.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 = VCompleter(smart_completion, vspecial=self.vspecial)
self.completer = completer
self._completer_lock = threading.Lock()
self.register_special_commands()
self.cli = None
def register_special_commands(self):
self.vspecial.register(self.change_db, '\\c',
'\\c[onnect] [DBNAME]',
'Connect to a new database',
aliases=('use', '\\connect', 'USE'))
self.vspecial.register(self.refresh_completions, '\\#', '\\#',
'Refresh auto-completions', arg_type=NO_QUERY)
self.vspecial.register(self.refresh_completions, '\\refresh',
'\\refresh', 'Refresh auto-completions',
arg_type=NO_QUERY)
def change_db(self, pattern, **_):
if pattern:
db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern
self.vexecute.connect(database=db)
else:
self.vexecute.connect()
yield (None, None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.vexecute.dbname, self.vexecute.user), True)
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('vcli')
root_logger.addHandler(handler)
root_logger.setLevel(level_map[log_level.upper()])
root_logger.debug('Initializing vcli logging.')
root_logger.debug('Log file %r.', log_file)
def connect_uri(self, uri):
uri = urlparse(uri)
database = uri.path[1:] # ignore the leading fwd slash
host = uri.hostname or 'localhost'
user = uri.username or getpass.getuser()
port = uri.port or 5433
password = uri.password or ''
self.connect(database, host, user, port, password)
def connect(self, database, host, user, port, passwd):
# Connect to the database
try:
self.vexecute = VExecute(database, user, passwd, host, port)
except errors.DatabaseError as e: # Connection can fail
self.logger.debug('Database connection failed: %r.', e)
self.logger.error("traceback: %r", traceback.format_exc())
error_msg = str(e) or type(e).__name__
click.secho(error_msg, err=True, fg='red')
exit(1)
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):
vexecute = self.vexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
completer = self.completer
self.refresh_completions()
def set_vi_mode(value):
self.vi_mode = value
key_binding_manager = vcli_bindings(
get_vi_mode_enabled=lambda: self.vi_mode,
set_vi_mode_enabled=set_vi_mode)
click.secho('Version: %s' % __version__)
def prompt_tokens(cli):
return [(Token.Prompt, '%s=> ' % vexecute.dbname)]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.vi_mode,
self.completion_refresher.is_refreshing)
input_processors = [
# Highlight matching brackets while editing.
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
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=input_processors)
history_file = self.config['main']['history_file']
with self._completer_lock:
buf = VBuffer(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 vexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the vexecute.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 vexecute.run(). Which causes
# finally clause to fail.
res = []
start = time()
# Run the query.
res = vexecute.run(document.text, self.vspecial)
duration = time() - start
file_output = None
stdout_output = []
total = 0
for title, cur, headers, status, force_stdout in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and
cur and cur.rowcount > threshold):
click.secho('The result set has more than %s rows.'
% threshold, fg='red')
if not click.confirm('Do you want to continue?'):
click.secho("Aborted!", err=True, fg='red')
break
formatted = format_output(title, cur, headers, status,
self.table_format,
self.vspecial.expanded_output,
self.vspecial.aligned,
self.vspecial.show_header)
if self.vspecial.output is not sys.stdout:
file_output = self.vspecial.output
if force_stdout or not file_output:
output = stdout_output
else:
output = file_output
write_output(output, formatted)
if hasattr(cur, 'rowcount'):
if self.vspecial.show_header:
if cur.rowcount == 1:
write_output(output, '(1 row)')
elif headers:
write_output(output, '(%d rows)' % cur.rowcount)
if document.text.startswith('\\') and cur.rowcount == 0:
stdout_output = ['No matching relations found.']
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
vexecute.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 errors.ConnectionError as e:
reconnect = True
if ('Connection is closed' in utf8tounicode(e.args[0])):
reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
show_default=False, type=bool, default=True)
if reconnect:
try:
vexecute.connect()
click.secho('Reconnected!\nTry the command again.', fg='green')
except errors.DatabaseError 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:
successful = True
if stdout_output:
output = '\n'.join(stdout_output)
try:
click.echo_via_pager(output)
except KeyboardInterrupt:
pass
if file_output:
try:
file_output.flush()
except KeyboardInterrupt:
pass
if self.vspecial.timing_enabled:
print('Time: command: %0.03fs, total: %0.03fs' % (duration, total))
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions(need_completion_reset(document.text))
# 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(vexecute.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'] = '-RXF'
return less_opts
def refresh_completions(self, reset=False):
if reset:
with self._completer_lock:
self.completer.reset_completions()
self.completion_refresher.refresh(self.vexecute, self.vspecial,
self._on_completions_refreshed)
return [(None, None, None,
'Auto-completion refresh started in the background.', True)]
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 run_cli(self):
sqlexecute = self.sqlexecute
logger = self.logger
original_less_opts = self.adjust_less_opts()
self.initialize_completions()
completer = self.completer
def set_key_bindings(value):
if value not in ('emacs', 'vi'):
value = 'emacs'
self.key_bindings = value
project_root = os.path.dirname(PACKAGE_ROOT)
author_file = os.path.join(project_root, 'AUTHORS')
sponsor_file = os.path.join(project_root, 'SPONSORS')
key_binding_manager = mycli_bindings(get_key_bindings=lambda: self.key_bindings,
set_key_bindings=set_key_bindings)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/mycli')
print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
print('Home: http://mycli.net')
print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file]))
def prompt_tokens(cli):
return [(Token.Prompt, self.get_prompt(self.prompt_format))]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.key_bindings)
layout = create_default_layout(lexer=MyCliLexer,
reserve_space_for_menu=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
])
buf = CLIBuffer(always_multiline=self.multi_line, completer=completer,
history=FileHistory(os.path.expanduser('~/.mycli-history')),
complete_while_typing=Always())
application = Application(style=style_factory(self.syntax_style),
layout=layout, buffer=buf,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION)
cli = CommandLineInterface(application=application, eventloop=create_eventloop())
try:
while True:
document = cli.run()
special.set_expanded_output(False)
# The reason we check here instead of inside the sqlexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the
# sqlexecute.run() statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
continue
destroy = confirm_destructive_query(document.text)
if destroy is None:
pass # Query was not destructive. Nothing to do here.
elif destroy is True:
self.output('Your call!')
else:
self.output('Wise choice!')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
if self.logfile:
self.logfile.write('\n# %s\n' % datetime.now())
self.logfile.write(document.text)
self.logfile.write('\n')
successful = False
start = time()
res = sqlexecute.run(document.text)
duration = time() - start
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and
cur and cur.rowcount > threshold):
self.output('The result set has more than %s rows.'
% threshold, fg='red')
if not click.confirm('Do you want to continue?'):
self.output("Aborted!", err=True, fg='red')
break
output.extend(format_output(title, cur, headers,
status, self.table_format))
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except KeyboardInterrupt:
# Restart connection to the database
sqlexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
self.output("cancelled query", err=True, fg='red')
except NotImplementedError:
self.output('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
logger.debug("Exception: %r", e)
reconnect = True
if (e.args[0] in (2003, 2006, 2013)):
reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
show_default=False, type=bool, default=True)
if reconnect:
logger.debug('Attempting to reconnect.')
try:
sqlexecute.connect()
logger.debug('Reconnected successfully.')
self.output('Reconnected!\nTry the command again.', fg='green')
except OperationalError as e:
logger.debug('Reconnect failed. e: %r', e)
self.output(str(e), err=True, fg='red')
continue # If reconnection failed, don't proceed further.
else: # If user chooses not to reconnect, don't proceed further.
continue
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
else:
self.output_via_pager('\n'.join(output))
if special.is_timing_enabled():
self.output('Command Time: %0.03fs' % duration)
self.output('Format Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_dynamic_completions()
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
self.output('Goodbye!')
finally: # Reset the less opts back to original.
logger.debug('Restoring env var LESS to %r.', original_less_opts)
os.environ['LESS'] = original_less_opts
os.environ['PAGER'] = special.get_original_pager()
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 = []
start = time()
# Run the query.
res = pgexecute.run(document.text, self.pgspecial)
duration = time() - start
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
threshold = 1000
if (is_select(status) and cur
and cur.rowcount > threshold):
click.secho(
'The result set has more than %s rows.' %
threshold,
fg='red')
if not click.confirm('Do you want to continue?'):
click.secho("Aborted!", err=True, fg='red')
break
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('Command Time: %0.03fs' % duration)
print('Format Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions()
# Refresh search_path to set default schema.
if need_search_path_refresh(document.text):
logger.debug('Refreshing search path')
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)
class MyCli(object):
default_prompt = '\\t \\u@\\h:\\d> '
defaults_suffix = None
# In order of being loaded. Files lower in list override earlier ones.
cnf_files = [
'/etc/my.cnf',
'/etc/mysql/my.cnf',
'/usr/local/etc/my.cnf',
'~/.my.cnf'
]
system_config_files = [
'/etc/myclirc',
]
default_config_file = os.path.join(PACKAGE_ROOT, 'myclirc')
user_config_file = '~/.myclirc'
def __init__(self, sqlexecute=None, prompt=None,
logfile=None, defaults_suffix=None, defaults_file=None,
login_path=None, auto_vertical_output=False):
self.sqlexecute = sqlexecute
self.logfile = logfile
self.defaults_suffix = defaults_suffix
self.login_path = login_path
self.auto_vertical_output = auto_vertical_output
# self.cnf_files is a class variable that stores the list of mysql
# config files to read in at launch.
# If defaults_file is specified then override the class variable with
# defaults_file.
if defaults_file:
self.cnf_files = [defaults_file]
# Load config.
config_files = ([self.default_config_file] + self.system_config_files +
[self.user_config_file])
c = self.config = read_config_files(config_files)
self.multi_line = c['main'].as_bool('multi_line')
self.destructive_warning = c['main'].as_bool('destructive_warning')
self.key_bindings = c['main']['key_bindings']
special.set_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')
# Write user config if system config wasn't the last config loaded.
if c.filename not in self.system_config_files:
write_default_config(self.default_config_file, self.user_config_file)
# audit log
if self.logfile is None and 'audit_log' in c['main']:
try:
self.logfile = open(os.path.expanduser(c['main']['audit_log']), 'a')
except (IOError, OSError) as e:
self.output('Error: Unable to open the audit log file. Your queries will not be logged.', err=True, fg='red')
self.logfile = False
self.completion_refresher = CompletionRefresher()
self.logger = logging.getLogger(__name__)
self.initialize_logging()
prompt_cnf = self.read_my_cnf_files(self.cnf_files, ['prompt'])['prompt']
self.prompt_format = prompt or prompt_cnf or c['main']['prompt'] or \
self.default_prompt
self.query_history = []
# Initialize completer.
smart_completion = c['main'].as_bool('smart_completion')
self.completer = SQLCompleter(smart_completion)
self._completer_lock = threading.Lock()
# Register custom special commands.
self.register_special_commands()
# Load .mylogin.cnf if it exists.
mylogin_cnf_path = get_mylogin_cnf_path()
if mylogin_cnf_path:
try:
mylogin_cnf = open_mylogin_cnf(mylogin_cnf_path)
if mylogin_cnf_path and mylogin_cnf:
# .mylogin.cnf gets read last, even if defaults_file is specified.
self.cnf_files.append(mylogin_cnf)
elif mylogin_cnf_path and not mylogin_cnf:
# There was an error reading the login path file.
print('Error: Unable to read login path file.')
except CryptoError:
click.secho('Warning: .mylogin.cnf was not read: pycrypto '
'module is not available.')
self.cli = None
def register_special_commands(self):
special.register_special_command(self.change_db, 'use',
'\\u', 'Change to a new database.', aliases=('\\u',))
special.register_special_command(self.change_db, 'connect',
'\\r', 'Reconnect to the database. Optional database argument.',
aliases=('\\r', ), case_sensitive=True)
special.register_special_command(self.refresh_completions, 'rehash',
'\\#', 'Refresh auto-completions.', arg_type=NO_QUERY, aliases=('\\#',))
special.register_special_command(self.change_table_format, 'tableformat',
'\\T', 'Change Table Type.', aliases=('\\T',), case_sensitive=True)
special.register_special_command(self.execute_from_file, 'source', '\\. filename',
'Execute commands from file.', aliases=('\\.',))
special.register_special_command(self.change_prompt_format, 'prompt',
'\\R', 'Change prompt format.', aliases=('\\R',), case_sensitive=True)
def change_table_format(self, arg, **_):
if not arg in table_formats():
msg = "Table type %s not yet implemented. Allowed types:" % arg
for table_type in table_formats():
msg += "\n\t%s" % table_type
yield (None, None, None, msg)
else:
self.table_format = arg
yield (None, None, None, "Changed table Type to %s" % self.table_format)
def change_db(self, arg, **_):
if arg is None:
self.sqlexecute.connect()
else:
self.sqlexecute.connect(database=arg)
yield (None, None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.sqlexecute.dbname, self.sqlexecute.user))
def execute_from_file(self, arg, **_):
if not arg:
message = 'Missing required argument, filename.'
return [(None, None, None, message)]
try:
with open(os.path.expanduser(arg), encoding='utf-8') as f:
query = f.read()
except IOError as e:
return [(None, None, None, str(e))]
return self.sqlexecute.run(query)
def change_prompt_format(self, arg, **_):
"""
Change the prompt format.
"""
if not arg:
message = 'Missing required argument, format.'
return [(None, None, None, message)]
self.prompt_format = self.get_prompt(arg)
return [(None, None, None, "Changed prompt format to %s" % arg)]
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('mycli')
root_logger.addHandler(handler)
root_logger.setLevel(level_map[log_level.upper()])
logging.captureWarnings(True)
root_logger.debug('Initializing mycli logging.')
root_logger.debug('Log file %r.', log_file)
def connect_uri(self, uri, local_infile=None):
uri = urlparse(uri)
database = uri.path[1:] # ignore the leading fwd slash
self.connect(database, uri.username, uri.password, uri.hostname,
uri.port, local_infile=local_infile)
def read_my_cnf_files(self, files, keys):
"""
Reads a list of config files and merges them. The last one will win.
:param files: list of files to read
:param keys: list of keys to retrieve
:returns: tuple, with None for missing keys.
"""
cnf = read_config_files(files)
sections = ['client']
if self.login_path and self.login_path != 'client':
sections.append(self.login_path)
if self.defaults_suffix:
sections.extend([sect + self.defaults_suffix for sect in sections])
def get(key):
result = None
for sect in cnf:
if sect in sections and key in cnf[sect]:
result = cnf[sect][key]
return result
return dict([(x, get(x)) for x in keys])
def connect(self, database='', user='', passwd='', host='', port='',
socket='', charset='', local_infile=''):
cnf = {'database': None,
'user': None,
'password': None,
'host': None,
'port': None,
'socket': None,
'default-character-set': None,
'local-infile': None,
'loose-local-infile': None}
cnf = self.read_my_cnf_files(self.cnf_files, cnf.keys())
# Fall back to config values only if user did not specify a value.
database = database or cnf['database']
if port or host:
socket = ''
else:
socket = socket or cnf['socket']
user = user or cnf['user'] or os.getenv('USER')
host = host or cnf['host'] or 'localhost'
port = port or cnf['port'] or 3306
try:
port = int(port)
except ValueError as e:
self.output("Error: Invalid port number: '{0}'.".format(port),
err=True, fg='red')
exit(1)
passwd = passwd or cnf['password']
charset = charset or cnf['default-character-set'] or 'utf8'
# Favor whichever local_infile option is set.
for local_infile_option in (local_infile, cnf['local-infile'],
cnf['loose-local-infile'], False):
try:
local_infile = str_to_bool(local_infile_option)
break
except (TypeError, ValueError):
pass
# Connect to the database.
try:
try:
sqlexecute = SQLExecute(database, user, passwd, host, port,
socket, charset, local_infile)
except OperationalError as e:
if ('Access denied for user' in e.args[1]):
passwd = click.prompt('Password', hide_input=True,
show_default=False, type=str)
sqlexecute = SQLExecute(database, user, passwd, host, port,
socket, charset, local_infile)
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())
self.output(str(e), err=True, fg='red')
exit(1)
self.sqlexecute = sqlexecute
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):
sqlexecute = self.sqlexecute
logger = self.logger
self.configure_pager()
self.refresh_completions()
def set_key_bindings(value):
if value not in ('emacs', 'vi'):
value = 'emacs'
self.key_bindings = value
project_root = os.path.dirname(PACKAGE_ROOT)
author_file = os.path.join(project_root, 'AUTHORS')
sponsor_file = os.path.join(project_root, 'SPONSORS')
key_binding_manager = mycli_bindings(get_key_bindings=lambda: self.key_bindings,
set_key_bindings=set_key_bindings)
print('Version:', __version__)
print('Chat: https://gitter.im/dbcli/mycli')
print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
print('Home: http://mycli.net')
print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file]))
def prompt_tokens(cli):
return [(Token.Prompt, self.get_prompt(self.prompt_format))]
get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.key_bindings,
self.completion_refresher.is_refreshing)
layout = create_prompt_layout(lexer=MyCliLexer,
multiline=True,
get_prompt_tokens=prompt_tokens,
get_bottom_toolbar_tokens=get_toolbar_tokens,
display_completions_in_columns=self.wider_completion_menu,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
])
with self._completer_lock:
buf = CLIBuffer(always_multiline=self.multi_line, completer=self.completer,
history=FileHistory(os.path.expanduser('~/.mycli-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)
self.cli = CommandLineInterface(application=application,
eventloop=create_eventloop())
try:
while True:
document = self.cli.run()
special.set_expanded_output(False)
# The reason we check here instead of inside the sqlexecute is
# because we want to raise the Exit exception which will be
# caught by the try/except block that wraps the
# sqlexecute.run() statement.
if quit_command(document.text):
raise EOFError
try:
document = self.handle_editor_command(self.cli, document)
except RuntimeError as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
continue
if self.destructive_warning:
destroy = confirm_destructive_query(document.text)
if destroy is None:
pass # Query was not destructive. Nothing to do here.
elif destroy is True:
self.output('Your call!')
else:
self.output('Wise choice!')
continue
# Keep track of whether or not the query is mutating. In case
# of a multi-statement query, the overall query is considered
# mutating if any one of the component statements is mutating
mutating = False
try:
logger.debug('sql: %r', document.text)
if self.logfile:
self.logfile.write('\n# %s\n' % datetime.now())
self.logfile.write(document.text)
self.logfile.write('\n')
successful = False
start = time()
res = sqlexecute.run(document.text)
successful = True
output = []
total = 0
for title, cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
threshold = 1000
if (is_select(status) and
cur and cur.rowcount > threshold):
self.output('The result set has more than %s rows.'
% threshold, fg='red')
if not click.confirm('Do you want to continue?'):
self.output("Aborted!", err=True, fg='red')
break
if self.auto_vertical_output:
max_width = self.cli.output.get_size().columns
else:
max_width = None
formatted = format_output(title, cur, headers,
status, self.table_format,
special.is_expanded_output(), max_width)
output.extend(formatted)
end = time()
total += end - start
mutating = mutating or is_mutating(status)
except UnicodeDecodeError as e:
import pymysql
if pymysql.VERSION < ('0', '6', '7'):
message = ('You are running an older version of pymysql.\n'
'Please upgrade to 0.6.7 or above to view binary data.\n'
'Try \'pip install -U pymysql\'.')
self.output(message)
else:
raise e
except KeyboardInterrupt:
# Restart connection to the database
sqlexecute.connect()
logger.debug("cancelled query, sql: %r", document.text)
self.output("cancelled query", err=True, fg='red')
except NotImplementedError:
self.output('Not Yet Implemented.', fg="yellow")
except OperationalError as e:
logger.debug("Exception: %r", e)
reconnect = True
if (e.args[0] in (2003, 2006, 2013)):
reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
show_default=False, type=bool, default=True)
if reconnect:
logger.debug('Attempting to reconnect.')
try:
sqlexecute.connect()
logger.debug('Reconnected successfully.')
self.output('Reconnected!\nTry the command again.', fg='green')
except OperationalError as e:
logger.debug('Reconnect failed. e: %r', e)
self.output(str(e), err=True, fg='red')
continue # If reconnection failed, don't proceed further.
else: # If user chooses not to reconnect, don't proceed further.
continue
else:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
except Exception as e:
logger.error("sql: %r, error: %r", document.text, e)
logger.error("traceback: %r", traceback.format_exc())
self.output(str(e), err=True, fg='red')
else:
try:
if special.is_pager_enabled():
self.output_via_pager('\n'.join(output))
else:
self.output('\n'.join(output))
except KeyboardInterrupt:
pass
if special.is_timing_enabled():
self.output('Time: %0.03fs' % total)
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):
self.refresh_completions(
reset=need_completion_reset(document.text))
finally:
if self.logfile is False:
self.output("Warning: This query was not logged.", err=True, fg='red')
query = Query(document.text, successful, mutating)
self.query_history.append(query)
except EOFError:
self.output('Goodbye!')
def output(self, text, **kwargs):
if self.logfile:
self.logfile.write(utf8tounicode(text))
self.logfile.write('\n')
click.secho(text, **kwargs)
def output_via_pager(self, text):
if self.logfile:
self.logfile.write(text)
self.logfile.write('\n')
click.echo_via_pager(text)
def configure_pager(self):
# Provide sane defaults for less.
os.environ['LESS'] = '-SRXF'
cnf = self.read_my_cnf_files(self.cnf_files, ['pager', 'skip-pager'])
if cnf['pager']:
special.set_pager(cnf['pager'])
if cnf['skip-pager']:
special.disable_pager()
def refresh_completions(self, reset=False):
if reset:
with self._completer_lock:
self.completer.reset_completions()
self.completion_refresher.refresh(self.sqlexecute,
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 mycli 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):
sqlexecute = self.sqlexecute
string = string.replace('\\u', sqlexecute.user or '(none)')
string = string.replace('\\h', sqlexecute.host or '(none)')
string = string.replace('\\d', sqlexecute.dbname or '(none)')
string = string.replace('\\t', sqlexecute.server_type()[0] or 'mycli')
string = string.replace('\\n', "\n")
return string
def loop(cmd, history_file):
from prompt_toolkit import CommandLineInterface, AbortAction, Application
from prompt_toolkit.interface import AcceptAction
from prompt_toolkit.enums import DEFAULT_BUFFER
from prompt_toolkit.layout.processors import (
HighlightMatchingBracketProcessor,
ConditionalProcessor
)
from prompt_toolkit.key_binding.manager import KeyBindingManager
from prompt_toolkit.shortcuts import (create_default_layout,
create_default_output,
create_eventloop)
key_binding_manager = KeyBindingManager(
enable_search=True,
enable_abort_and_exit_bindings=True,
enable_vi_mode=Condition(lambda cli: _enable_vi_mode()))
layout = create_default_layout(
message=u'cr> ',
multiline=True,
lexer=SqlLexer,
extra_input_processors=[
ConditionalProcessor(
processor=HighlightMatchingBracketProcessor(chars='[](){}'),
filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
]
)
cli_buffer = CrashBuffer(
history=TruncatedFileHistory(history_file, max_length=MAX_HISTORY_LENGTH),
accept_action=AcceptAction.RETURN_DOCUMENT,
completer=SQLCompleter(cmd.connection, cmd.lines),
complete_while_typing=Always()
)
application = Application(
layout=layout,
style=CrateStyle,
buffer=cli_buffer,
key_bindings_registry=key_binding_manager.registry,
on_exit=AbortAction.RAISE_EXCEPTION,
on_abort=AbortAction.RETRY,
)
eventloop = create_eventloop()
output = create_default_output()
cli = CommandLineInterface(
application=application,
eventloop=eventloop,
output=output
)
def get_num_columns_override():
return output.get_size().columns
cmd.get_num_columns = get_num_columns_override
while True:
try:
doc = cli.run()
if doc:
cmd.process(doc.text)
except EOFError:
cmd.logger.warn(u'Bye!')
return
