def main():
    eventloop = create_eventloop()

    cli = CommandLineInterface(layout=Window(BufferControl(input_processors=[ClockPrompt()])),
                               eventloop=eventloop)
    done = [False]  # Non local

    def on_read_start():
        """
        This function is called when we start reading at the input.
        (Actually the start of the read-input event loop.)
        """
        # Following function should be run in the background.
        # We do it by using an executor thread from the `CommandLineInterface`
        # instance.
        def run():
            # Send every second a redraw request.
            while not done[0]:
                time.sleep(1)
                cli.request_redraw()

        cli.eventloop.run_in_executor(run)

    def on_read_end():
        done[0] = True

    cli.onReadInputStart += on_read_start
    cli.onReadInputEnd += on_read_end

    code_obj = cli.read_input()
    print('You said: %s' % code_obj.text)

    eventloop.close()
Esempio n. 2
0
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 main():
    cli = CommandLineInterface(layout=Layout(
        before_input=DefaultPrompt('Enter HTML: '),
        lexer=HtmlLexer))

    html_code_obj = cli.read_input()
    print('You said: ' + html_code_obj.text)
def main():
    layout = Layout(
        left_margin=LeftMarginWithLineNumbers(),
        before_input=DefaultPrompt(text='Before input >> '),
        after_input=Prompt(' << after input'),
        top_toolbars=[
            TextToolbar('This is a top toolbar\n(second line)\n(third line)', token=Token.TopToolbar1, height=3),
            TextToolbar('This is another top toolbar', token=Token.TopToolbar2),
        ],
        bottom_toolbars=[
            ArgToolbar(),
            SearchToolbar(),
            CompletionsToolbar(),
            TextToolbar('This is a bottom toolbar', token=Token.BottomToolbar1),
            TextToolbar('This is a multiline bottom toolbar\n(second line)\n(third line)', token=Token.BottomToolbar2, height=3),
        ],
        show_tildes=True,
        menus=[CompletionsMenu()])

    cli = CommandLineInterface(layout=layout,
                               style=TestStyle,
                               line=Line(is_multiline=True, completer=TestCompleter()))

    code_obj = cli.read_input(initial_value=lipsum)
    print('You said: ' + code_obj.text)
def main():
    cli = CommandLineInterface(line=Line(history=FileHistory('.example-history-file')))

    try:
        while True:
            document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
            print('You said: ' + document.text)
    except Exit:
        pass
Esempio n. 6
0
    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!")
Esempio n. 7
0
def main():
    layout = Layout(before_input=DefaultPrompt('> '), lexer=HyLexer, menus=[CompletionMenu()])
    line = Line(completer=HyCompleter())
    cli = CommandLineInterface(layout=layout, line=line)
    try:
        while True:
            code_obj = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
            print 'You entered:', code_obj.text
    except Exit:
        print 'GoodBye!'
def main():
    cli = CommandLineInterface(style=AnimalStyle,
                   layout=Layout(before_input=DefaultPrompt('Give some animals: '),
                                 menus=[CompletionMenu()]),
                   line=Line(completer=AnimalCompleter()),
                   create_async_autocompleters=True,
            )

    print('Press tab to complete')
    code_obj = cli.read_input()
    print('You said: ' + code_obj.text)
Esempio n. 9
0
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
Esempio n. 10
0
    def _get_input(self):
        """
        Read PDB input. Return input text.
        """
        g = self._create_grammar()

        cli = CommandLineInterface(
            layout=Layout(
                before_input=PdbPrompt(self._get_current_pdb_commands()),
                show_tildes=True,
                min_height=15,
                lexer=GrammarLexer(
                    g,
                    tokens={'pdb_command': Token.PdbCommand},
                    lexers={'python_code': PythonLexer}
                ),
                after_input=CompletionHint(),
                menus=[CompletionsMenu()],
                top_toolbars=[],
                bottom_toolbars=[
                    SystemToolbar(),
                    ArgToolbar(),
                    SearchToolbar(),
                    SourceCodeToolbar(weakref.ref(self)),
                    ValidationToolbar(),
                    ShortcutsToolbar(),
                    PdbStatusToolbar(weakref.ref(self)),
                ]),
            line=Line(
                completer=GrammarCompleter(g, completers={
                    'enabled_breakpoint': BreakPointListCompleter(only_enabled=True),
                    'disabled_breakpoint': BreakPointListCompleter(only_disabled=True),
                    'alias_name': AliasCompleter(self),
                    'python_code': PythonCompleter(lambda: self.curframe.f_globals, lambda: self.curframe.f_locals),
                    'breakpoint': BreakPointListCompleter(),
                    'pdb_command': PdbCommandsCompleter(self),
                    'python_file': PythonFileCompleter(),
                    'python_function': PythonFunctionCompleter(self),
                }),
                history=self._cli_history,
                validator=GrammarValidator(g, {
                    'python_code': PythonValidator()
                }),
            ),
            key_binding_factories=[emacs_bindings, custom_pdb_key_bindings],
            style=PdbStyle)

        try:
            return cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION).text
        except Exit:
            # Turn Control-D key press into a 'quit' command.
            return 'q'
def main(database):
    connection = sqlite3.connect(database)
    layout = Layout(before_input=DefaultPrompt('> '),
                    lexer=SqlLexer, menus=[CompletionMenu()])
    line = Line(completer=SqlCompleter())
    cli = CommandLineInterface(style=DocumentStyle, layout=layout, line=line)
    try:
        while True:
            document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
            with connection:
                messages = connection.execute(document.text)
                for message in messages:
                    print message
    except Exit:
        print 'GoodBye!'
Esempio n. 12
0
 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)
Esempio n. 13
0
def main():
    layout = Layout(before_input=DefaultPrompt("> "), menus=[CompletionMenu()])
    line = Line(RESTCompleter())
    cli = CommandLineInterface(style=DocumentStyle, layout=layout, line=line)
    try:
        while True:
            document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)

            input_args = document.text.split(" ")

            if len(input_args) < 2:
                raise AssertionError("Must provide at least a method and a url")

            response = process_request(input_args)

            print "Response:", response.json()

    except Exit:
        print "GoodBye!"
Esempio n. 14
0
    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)
Esempio n. 15
0
def cli(database, user, password, host, port):

    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.
    write_default_config(default_config, '~/.pgclirc')

    # Load config.
    config = load_config('~/.pgclirc')

    # Connect to the database.
    try:
        pgexecute = PGExecute(database, user, password, host, port)
    except Exception as e:
        click.secho(e.message, err=True, fg='red')
        exit(1)
    layout = Layout(before_input=DefaultPrompt('%s> ' % database),
            menus=[CompletionsMenu()],
            lexer=SqlLexer)
    completer = PGCompleter(config.getboolean('main', 'smart_completion'))
    completer.extend_special_commands(pgexecute.special_commands.keys())
    completer.extend_table_names(pgexecute.tables())
    completer.extend_column_names(pgexecute.all_columns())
    line = Line(completer=completer,
            history=FileHistory(os.path.expanduser('~/.pgcli-history')))
    cli = CommandLineInterface(style=PGStyle, layout=layout, line=line)

    try:
        while True:
            document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
            try:
                rows, headers, status = pgexecute.run(document.text)
                if rows:
                    print(tabulate(rows, headers, tablefmt='psql'))
                print(status)
            except Exception as e:
                click.secho(e.message, err=True, fg='red')
    except Exit:
        print ('GoodBye!')
Esempio n. 16
0
def get_input(message, raise_exception_on_abort=False, multiline=False,
        is_password=False, vi_mode=False):
    """
    Replacement for `raw_input`.
    Ask for input, return the answer.
    This returns `None` when Ctrl-D was pressed.
    """
    layout = Layout(
        before_input=DefaultPrompt(message),
        input_processors=([PasswordProcessor()] if is_password else []))

    cli = CommandLineInterface(
        layout=layout,
        line=Line(is_multiline=multiline),
        key_binding_factories=[(vi_bindings if vi_mode else emacs_bindings)])

    on_abort = AbortAction.RAISE_EXCEPTION if raise_exception_on_abort else AbortAction.RETURN_NONE
    code = cli.read_input(on_abort=on_abort, on_exit=AbortAction.IGNORE)

    if code:
        return code.text
Esempio n. 17
0
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!')
Esempio n. 18
0
def main():
    cli = CommandLineInterface(layout=Layout(before_input=ClockPrompt()))

    def on_read_start():
        """
        This function is called when we start reading at the input.
        (Actually the start of the read-input event loop.)
        """
        # Following function should be run in the background.
        # We do it by using an executor thread from the `CommandLineInterface`
        # instance.
        def run():
            # Send every second a redraw request.
            while cli.is_reading_input:
                time.sleep(1)
                cli.request_redraw()

        cli.run_in_executor(run)
    cli.onReadInputStart += on_read_start

    code_obj = cli.read_input()
    print('You said: ' + code_obj.text)
def interactive_shell():
    """
    Coroutine that shows the interactive command line.
    """
    # Create interface. (style/layout is only for demonstration.)
    cli = CommandLineInterface(
        layout=Layout(before_input=DefaultPrompt(text='Say something inside the event loop: ')),
        style=TestStyle)

    # Patch stdout in something that will always print *above* the prompt when
    # something is written to stdout.
    sys.stdout = cli.stdout_proxy()

    # Run echo loop. Read text from stdin, and reply it back.
    while True:
        try:
            result = yield from cli.read_input_async(
                on_exit=AbortAction.RAISE_EXCEPTION,
                on_abort=AbortAction.RAISE_EXCEPTION)
            print('You said: "%s"' % result.text)
        except (Exit, Abort):
            loop.stop()
            print('Qutting event loop. Bye.')
            return
Esempio n. 20
0
 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 main():
    manager = KeyBindingManager(enable_system_prompt=True)

    D = LayoutDimension
    layout = HSplit([
        VSplit([
            Window(width=D(min=15, max=30, preferred=30),
                   content=FillControl('a', token=Token.A)),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            Window(content=TokenListControl.static([(Token.HelloWorld, lipsum)])),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            Window(content=BufferControl(lexer=PythonLexer,
                                         show_line_numbers=Always(),
                                         input_processors=[
                                                DefaultPrompt('python> '),
                                                AfterInput.static(' <python', token=Token.AfterInput),
                                         ]),
            ),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            HSplit([
                Window(width=D(max=40),
                      height=D.exact(4),
                      content=FillControl('b', token=Token.B)),
                Window(width=D(max=40),
                      content=FillControl('f', token=Token.F)),
                Window(width=D.exact(30),
                      height=D.exact(2),
                      content=FillControl('c', token=Token.C)),
            ]),
            #CompletionsMenu(),
        ]),
        Window(height=D.exact(1),
              content=FillControl('-', token=Token.Line)),
        Window(height=D.exact(3),
              content=FillControl('d', token=Token.D)),
        SystemToolbar(),
        ArgToolbar(),
        CompletionsToolbar(),
        SearchToolbar(),
    ])

    layout = FloatContainer(
        content=layout,
        floats=[
            Float(xcursor=True,
                  ycursor=True,
                  content=VSplit([
                      Window(width=D.exact(5),
                             content=FillControl('f', token=Token.F)),
                      CompletionsMenu(),
                  ])
            ),
        ]
    )

    eventloop = create_eventloop()
    cli = CommandLineInterface(eventloop=eventloop,
                               layout=layout,
                               style=TestStyle,
                               key_bindings_registry=manager.registry,
                               buffer=Buffer(is_multiline=Always(), completer=TestCompleter()))
    cli.read_input()
    eventloop.close()
Esempio n. 22
0
class PGCli(object):
    def __init__(self,
                 force_passwd_prompt=False,
                 never_passwd_prompt=False,
                 pgexecute=None,
                 pgclirc_file=None):

        self.force_passwd_prompt = force_passwd_prompt
        self.never_passwd_prompt = never_passwd_prompt
        self.pgexecute = pgexecute

        from pgcli import __file__ as package_root
        package_root = os.path.dirname(package_root)

        default_config = os.path.join(package_root, 'pgclirc')
        write_default_config(default_config, pgclirc_file)

        self.pgspecial = PGSpecial()

        # Load config.
        c = self.config = load_config(pgclirc_file, default_config)
        self.multi_line = c['main'].as_bool('multi_line')
        self.vi_mode = c['main'].as_bool('vi')
        self.pgspecial.timing_enabled = c['main'].as_bool('timing')
        self.table_format = c['main']['table_format']
        self.syntax_style = c['main']['syntax_style']
        self.cli_style = c['colors']
        self.wider_completion_menu = c['main'].as_bool('wider_completion_menu')
        self.completion_refresher = CompletionRefresher()

        self.logger = logging.getLogger(__name__)
        self.initialize_logging()

        self.query_history = []

        # Initialize completer
        smart_completion = c['main'].as_bool('smart_completion')
        completer = PGCompleter(smart_completion, pgspecial=self.pgspecial)
        self.completer = completer
        self._completer_lock = threading.Lock()
        self.register_special_commands()

        self.cli = None

    def register_special_commands(self):

        self.pgspecial.register(self.change_db,
                                '\\c',
                                '\\c[onnect] database_name',
                                'Change to a new database.',
                                aliases=('use', '\\connect', 'USE'))
        self.pgspecial.register(self.refresh_completions,
                                '\\#',
                                '\\#',
                                'Refresh auto-completions.',
                                arg_type=NO_QUERY)
        self.pgspecial.register(self.refresh_completions,
                                '\\refresh',
                                '\\refresh',
                                'Refresh auto-completions.',
                                arg_type=NO_QUERY)
        self.pgspecial.register(self.execute_from_file, '\\i', '\\i filename',
                                'Execute commands from file.')

    def change_db(self, pattern, **_):
        if pattern:
            db = pattern[1:-1] if pattern[0] == pattern[-1] == '"' else pattern
            self.pgexecute.connect(database=db)
        else:
            self.pgexecute.connect()

        yield (None, None, None, 'You are now connected to database "%s" as '
               'user "%s"' % (self.pgexecute.dbname, self.pgexecute.user))

    def execute_from_file(self, pattern, **_):
        if not pattern:
            message = '\\i: missing required argument'
            return [(None, None, None, message)]
        try:
            with open(os.path.expanduser(pattern), encoding='utf-8') as f:
                query = f.read()
        except IOError as e:
            return [(None, None, None, str(e))]

        return self.pgexecute.run(query, self.pgspecial)

    def initialize_logging(self):

        log_file = self.config['main']['log_file']
        log_level = self.config['main']['log_level']

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

        handler = logging.FileHandler(os.path.expanduser(log_file))

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

        handler.setFormatter(formatter)

        root_logger = logging.getLogger('pgcli')
        root_logger.addHandler(handler)
        root_logger.setLevel(level_map[log_level.upper()])

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

    def connect_dsn(self, dsn):
        self.connect(dsn=dsn)

    def connect_uri(self, uri):
        uri = urlparse(uri)
        database = uri.path[1:]  # ignore the leading fwd slash
        self.connect(database, uri.hostname, uri.username, uri.port,
                     uri.password)

    def connect(self,
                database='',
                host='',
                user='',
                port='',
                passwd='',
                dsn=''):
        # Connect to the database.

        if not user:
            user = getuser()

        if not database:
            database = user

        # If password prompt is not forced but no password is provided, try
        # getting it from environment variable.
        if not self.force_passwd_prompt and not passwd:
            passwd = os.environ.get('PGPASSWORD', '')

        # Prompt for a password immediately if requested via the -W flag. This
        # avoids wasting time trying to connect to the database and catching a
        # no-password exception.
        # If we successfully parsed a password from a URI, there's no need to
        # prompt for it, even with the -W flag
        if self.force_passwd_prompt and not passwd:
            passwd = click.prompt('Password',
                                  hide_input=True,
                                  show_default=False,
                                  type=str)

        # Prompt for a password after 1st attempt to connect without a password
        # fails. Don't prompt if the -w flag is supplied
        auto_passwd_prompt = not passwd and not self.never_passwd_prompt

        # Attempt to connect to the database.
        # Note that passwd may be empty on the first attempt. If connection
        # fails because of a missing password, but we're allowed to prompt for
        # a password (no -w flag), prompt for a passwd and try again.
        try:
            try:
                pgexecute = PGExecute(database, user, passwd, host, port, dsn)
            except OperationalError as e:
                if ('no password supplied' in utf8tounicode(e.args[0])
                        and auto_passwd_prompt):
                    passwd = click.prompt('Password',
                                          hide_input=True,
                                          show_default=False,
                                          type=str)
                    pgexecute = PGExecute(database, user, passwd, host, port,
                                          dsn)
                else:
                    raise e

        except Exception as e:  # Connecting to a database could fail.
            self.logger.debug('Database connection failed: %r.', e)
            self.logger.error("traceback: %r", traceback.format_exc())
            click.secho(str(e), err=True, fg='red')
            exit(1)

        self.pgexecute = pgexecute

    def handle_editor_command(self, cli, document):
        """
        Editor command is any query that is prefixed or suffixed
        by a '\e'. The reason for a while loop is because a user
        might edit a query multiple times.
        For eg:
        "select * from \e"<enter> to edit it in vim, then come
        back to the prompt with the edited query "select * from
        blah where q = 'abc'\e" to edit it again.
        :param cli: CommandLineInterface
        :param document: Document
        :return: Document
        """
        while special.editor_command(document.text):
            filename = special.get_filename(document.text)
            sql, message = special.open_external_editor(filename,
                                                        sql=document.text)
            if message:
                # Something went wrong. Raise an exception and bail.
                raise RuntimeError(message)
            cli.current_buffer.document = Document(sql,
                                                   cursor_position=len(sql))
            document = cli.run(False)
            continue
        return document

    def run_cli(self):
        pgexecute = self.pgexecute
        logger = self.logger
        original_less_opts = self.adjust_less_opts()

        self.refresh_completions()

        def set_vi_mode(value):
            self.vi_mode = value

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

        print('Version:', __version__)
        print('Chat: https://gitter.im/dbcli/pgcli')
        print('Mail: https://groups.google.com/forum/#!forum/pgcli')
        print('Home: http://pgcli.com')

        def prompt_tokens(cli):
            return [(Token.Prompt, '%s> ' % pgexecute.dbname)]

        get_toolbar_tokens = create_toolbar_tokens_func(
            lambda: self.vi_mode,
            lambda: self.completion_refresher.is_refreshing())

        layout = create_default_layout(
            lexer=PostgresLexer,
            reserve_space_for_menu=True,
            get_prompt_tokens=prompt_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            multiline=True,
            extra_input_processors=[
                # Highlight matching brackets while editing.
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
            ])
        history_file = self.config['main']['history_file']
        with self._completer_lock:
            buf = PGBuffer(always_multiline=self.multi_line,
                           completer=self.completer,
                           history=FileHistory(
                               os.path.expanduser(history_file)),
                           complete_while_typing=Always())

            application = Application(
                style=style_factory(self.syntax_style, self.cli_style),
                layout=layout,
                buffer=buf,
                key_bindings_registry=key_binding_manager.registry,
                on_exit=AbortAction.RAISE_EXCEPTION,
                ignore_case=True)
            self.cli = CommandLineInterface(application=application,
                                            eventloop=create_eventloop())

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

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

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

                # Keep track of whether or not the query is mutating. In case
                # of a multi-statement query, the overall query is considered
                # mutating if any one of the component statements is mutating
                mutating = False

                try:
                    logger.debug('sql: %r', document.text)
                    successful = False
                    # Initialized to [] because res might never get initialized
                    # if an exception occurs in pgexecute.run(). Which causes
                    # finally clause to fail.
                    res = []
                    # Run the query.
                    start = time()
                    res = pgexecute.run(document.text, self.pgspecial)
                    successful = True
                    output = []
                    total = 0
                    for title, cur, headers, status in res:
                        logger.debug("headers: %r", headers)
                        logger.debug("rows: %r", cur)
                        logger.debug("status: %r", status)
                        threshold = 1000
                        if (is_select(status) and cur
                                and cur.rowcount > threshold):
                            click.secho(
                                'The result set has more than %s rows.' %
                                threshold,
                                fg='red')
                            if not click.confirm('Do you want to continue?'):
                                click.secho("Aborted!", err=True, fg='red')
                                break

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

                        formatted = format_output(
                            title, cur, headers, status, self.table_format,
                            self.pgspecial.expanded_output, max_width)
                        output.extend(formatted)
                        end = time()
                        total += end - start
                        mutating = mutating or is_mutating(status)

                except KeyboardInterrupt:
                    # Restart connection to the database
                    pgexecute.connect()
                    logger.debug("cancelled query, sql: %r", document.text)
                    click.secho("cancelled query", err=True, fg='red')
                except NotImplementedError:
                    click.secho('Not Yet Implemented.', fg="yellow")
                except OperationalError as e:
                    reconnect = True
                    if ('server closed the connection'
                            in utf8tounicode(e.args[0])):
                        reconnect = click.prompt(
                            'Connection reset. Reconnect (Y/n)',
                            show_default=False,
                            type=bool,
                            default=True)
                        if reconnect:
                            try:
                                pgexecute.connect()
                                click.secho(
                                    'Reconnected!\nTry the command again.',
                                    fg='green')
                            except OperationalError as e:
                                click.secho(str(e), err=True, fg='red')
                    else:
                        logger.error("sql: %r, error: %r", document.text, e)
                        logger.error("traceback: %r", traceback.format_exc())
                        click.secho(str(e), err=True, fg='red')
                except Exception as e:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    click.secho(str(e), err=True, fg='red')
                else:
                    try:
                        click.echo_via_pager('\n'.join(output))
                    except KeyboardInterrupt:
                        pass
                    if self.pgspecial.timing_enabled:
                        print('Time: %0.03fs' % total)

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_completions()

                # Refresh search_path to set default schema.
                if need_search_path_refresh(document.text):
                    logger.debug('Refreshing search path')
                    with self._completer_lock:
                        self.completer.set_search_path(pgexecute.search_path())
                    logger.debug('Search path: %r', self.completer.search_path)

                query = Query(document.text, successful, mutating)
                self.query_history.append(query)

        except EOFError:
            print('Goodbye!')
        finally:  # Reset the less opts back to original.
            logger.debug('Restoring env var LESS to %r.', original_less_opts)
            os.environ['LESS'] = original_less_opts

    def adjust_less_opts(self):
        less_opts = os.environ.get('LESS', '')
        self.logger.debug('Original value for LESS env var: %r', less_opts)
        os.environ['LESS'] = '-SRXF'

        return less_opts

    def refresh_completions(self):
        self.completion_refresher.refresh(self.pgexecute, self.pgspecial,
                                          self._on_completions_refreshed)
        return [(None, None, None,
                 'Auto-completion refresh started in the background.')]

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

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

    def _swap_completer_objects(self, new_completer):
        """Swap the completer object in cli with the newly created completer.
        """
        with self._completer_lock:
            self.completer = new_completer
            # When pgcli is first launched we call refresh_completions before
            # instantiating the cli object. So it is necessary to check if cli
            # exists before trying the replace the completer object in cli.
            if self.cli:
                self.cli.current_buffer.completer = new_completer

    def get_completions(self, text, cursor_positition):
        with self._completer_lock:
            return self.completer.get_completions(
                Document(text=text, cursor_position=cursor_positition), None)
Esempio n. 23
0
class MyCli(object):

    default_prompt = '\\t \\u@\\h:\\d> '
    max_len_prompt = 45
    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')

    def __init__(self,
                 sqlexecute=None,
                 prompt=None,
                 logfile=None,
                 defaults_suffix=None,
                 defaults_file=None,
                 login_path=None,
                 auto_vertical_output=False,
                 warn=None,
                 myclirc="~/.myclirc"):
        self.sqlexecute = sqlexecute
        self.logfile = logfile
        self.defaults_suffix = defaults_suffix
        self.login_path = login_path

        # 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 +
                        [myclirc])
        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.formatter = TabularOutputFormatter(
            format_name=c['main']['table_format'])
        self.syntax_style = c['main']['syntax_style']
        self.less_chatty = c['main'].as_bool('less_chatty')
        self.cli_style = c['colors']
        self.wider_completion_menu = c['main'].as_bool('wider_completion_menu')
        c_dest_warning = c['main'].as_bool('destructive_warning')
        self.destructive_warning = c_dest_warning if warn is None else warn
        self.login_path_as_host = c['main'].as_bool('login_path_as_host')

        # read from cli argument or user config file
        self.auto_vertical_output = auto_vertical_output or \
                                c['main'].as_bool('auto_vertical_output')

        # 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, myclirc)

        # 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.echo(
                    '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.prompt_continuation_format = c['main']['prompt_continuation']

        self.query_history = []

        # Initialize completer.
        self.smart_completion = c['main'].as_bool('smart_completion')
        self.completer = SQLCompleter(
            self.smart_completion,
            supported_formats=self.formatter.supported_formats)
        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:
            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.')

        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 the table format used to output results.',
            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, **_):
        try:
            self.formatter.format_name = arg
            yield (None, None, None, 'Changed table format to {}'.format(arg))
        except ValueError:
            msg = 'Table format {} not recognized. Allowed formats:'.format(
                arg)
            for table_type in self.formatter.supported_formats:
                msg += "\n\t{}".format(table_type)
            yield (None, None, None, msg)

    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
        }

        # Disable logging if value is NONE by switching to a no-op handler
        # Set log level to a high value so it doesn't even waste cycles getting called.
        if log_level.upper() == "NONE":
            handler = logging.NullHandler()
            log_level = "CRITICAL"
        else:
            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, 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 {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.echo("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.echo(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
        """
        # FIXME: using application.pre_run_callables like this here is not the best solution.
        # It's internal api of prompt_toolkit that may change. This was added to fix
        # https://github.com/dbcli/pgcli/issues/668. We may find a better way to do it in the future.
        saved_callables = cli.application.pre_run_callables
        while special.editor_command(document.text):
            filename = special.get_filename(document.text)
            query = (special.get_editor_query(document.text)
                     or self.get_last_query())
            sql, message = special.open_external_editor(filename, sql=query)
            if message:
                # Something went wrong. Raise an exception and bail.
                raise RuntimeError(message)
            cli.current_buffer.document = Document(sql,
                                                   cursor_position=len(sql))
            cli.application.pre_run_callables = []
            document = cli.run()
            continue
        cli.application.pre_run_callables = saved_callables
        return document

    def run_cli(self):
        sqlexecute = self.sqlexecute
        logger = self.logger
        self.configure_pager()

        if self.smart_completion:
            self.refresh_completions()

        author_file = os.path.join(PACKAGE_ROOT, 'AUTHORS')
        sponsor_file = os.path.join(PACKAGE_ROOT, 'SPONSORS')

        key_binding_manager = mycli_bindings()

        if not self.less_chatty:
            print('Version:', __version__)
            print('Chat: https://gitter.im/dbcli/mycli')
            print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
            print('Home: http://mycli.net')
            print('Thanks to the contributor -',
                  thanks_picker([author_file, sponsor_file]))

        def prompt_tokens(cli):
            prompt = self.get_prompt(self.prompt_format)
            if self.prompt_format == self.default_prompt and len(
                    prompt) > self.max_len_prompt:
                prompt = self.get_prompt('\\d> ')
            return [(Token.Prompt, prompt)]

        def get_continuation_tokens(cli, width):
            continuation_prompt = self.get_prompt(
                self.prompt_continuation_format)
            return [(Token.Continuation, ' ' *
                     (width - len(continuation_prompt)) + continuation_prompt)]

        def one_iteration(document=None):
            if document is None:
                document = self.cli.run()

                special.set_expanded_output(False)

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

            if not document.text.strip():
                return

            if self.destructive_warning:
                destroy = confirm_destructive_query(document.text)
                if destroy is None:
                    pass  # Query was not destructive. Nothing to do here.
                elif destroy is True:
                    self.echo('Your call!')
                else:
                    self.echo('Wise choice!')
                    return

            # Keep track of whether or not the query is mutating. In case
            # of a multi-statement query, the overall query is considered
            # mutating if any one of the component statements is mutating
            mutating = False

            try:
                logger.debug('sql: %r', document.text)

                special.write_tee(
                    self.get_prompt(self.prompt_format) + document.text)
                if self.logfile:
                    self.logfile.write('\n# %s\n' % datetime.now())
                    self.logfile.write(document.text)
                    self.logfile.write('\n')

                successful = False
                start = time()
                res = sqlexecute.run(document.text)
                successful = True
                result_count = 0
                for title, cur, headers, status in res:
                    logger.debug("headers: %r", headers)
                    logger.debug("rows: %r", cur)
                    logger.debug("status: %r", status)
                    threshold = 1000
                    if (is_select(status) and cur
                            and cur.rowcount > threshold):
                        self.echo(
                            'The result set has more than {} rows.'.format(
                                threshold),
                            fg='red')
                        if not click.confirm('Do you want to continue?'):
                            self.echo("Aborted!", err=True, fg='red')
                            break

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

                    formatted = self.format_output(
                        title, cur, headers, special.is_expanded_output(),
                        max_width)

                    t = time() - start
                    try:
                        if result_count > 0:
                            self.echo('')
                        try:
                            self.output('\n'.join(formatted), status)
                        except KeyboardInterrupt:
                            pass
                        if special.is_timing_enabled():
                            self.echo('Time: %0.03fs' % t)
                    except KeyboardInterrupt:
                        pass

                    start = time()
                    result_count += 1
                    mutating = mutating or is_mutating(status)
                special.unset_once_if_written()
            except EOFError as e:
                raise e
            except KeyboardInterrupt:
                # get last connection id
                connection_id_to_kill = sqlexecute.connection_id
                logger.debug("connection id to kill: %r",
                             connection_id_to_kill)
                # Restart connection to the database
                sqlexecute.connect()
                try:
                    for title, cur, headers, status in sqlexecute.run(
                            'kill %s' % connection_id_to_kill):
                        status_str = str(status).lower()
                        if status_str.find('ok') > -1:
                            logger.debug(
                                "cancelled query, connection id: %r, sql: %r",
                                connection_id_to_kill, document.text)
                            self.echo("cancelled query", err=True, fg='red')
                except Exception as e:
                    self.echo(
                        'Encountered error while cancelling query: {}'.format(
                            e),
                        err=True,
                        fg='red')
            except NotImplementedError:
                self.echo('Not Yet Implemented.', fg="yellow")
            except OperationalError as e:
                logger.debug("Exception: %r", e)
                if (e.args[0] in (2003, 2006, 2013)):
                    logger.debug('Attempting to reconnect.')
                    self.echo('Reconnecting...', fg='yellow')
                    try:
                        sqlexecute.connect()
                        logger.debug('Reconnected successfully.')
                        one_iteration(document)
                        return  # OK to just return, cuz the recursion call runs to the end.
                    except OperationalError as e:
                        logger.debug('Reconnect failed. e: %r', e)
                        self.echo(str(e), err=True, fg='red')
                        # If reconnection failed, don't proceed further.
                        return
                else:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    self.echo(str(e), err=True, fg='red')
            except Exception as e:
                logger.error("sql: %r, error: %r", document.text, e)
                logger.error("traceback: %r", traceback.format_exc())
                self.echo(str(e), err=True, fg='red')
            else:
                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_completions(
                        reset=need_completion_reset(document.text))
            finally:
                if self.logfile is False:
                    self.echo("Warning: This query was not logged.",
                              err=True,
                              fg='red')
            query = Query(document.text, successful, mutating)
            self.query_history.append(query)

        get_toolbar_tokens = create_toolbar_tokens_func(
            self.completion_refresher.is_refreshing)

        layout = create_prompt_layout(
            lexer=MyCliLexer,
            multiline=True,
            get_prompt_tokens=prompt_tokens,
            get_continuation_tokens=get_continuation_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            extra_input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
            ],
            reserve_space_for_menu=self.get_reserved_space())
        with self._completer_lock:
            buf = CLIBuffer(always_multiline=self.multi_line,
                            completer=self.completer,
                            history=FileHistory(
                                os.path.expanduser(
                                    os.environ.get('MYCLI_HISTFILE',
                                                   '~/.mycli-history'))),
                            complete_while_typing=Always(),
                            accept_action=AcceptAction.RETURN_DOCUMENT)

            if self.key_bindings == 'vi':
                editing_mode = EditingMode.VI
            else:
                editing_mode = EditingMode.EMACS

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

        try:
            while True:
                one_iteration()
        except EOFError:
            special.close_tee()
            if not self.less_chatty:
                self.echo('Goodbye!')

    def log_output(self, output):
        """Log the output in the audit log, if it's enabled."""
        if self.logfile:
            self.logfile.write(utf8tounicode(output))
            self.logfile.write('\n')

    def echo(self, s, **kwargs):
        """Print a message to stdout.

        The message will be logged in the audit log, if enabled.

        All keyword arguments are passed to click.echo().

        """
        self.log_output(s)
        click.secho(s, **kwargs)

    def output_fits_on_screen(self, output, status=None):
        """Check if the given output fits on the screen."""
        size = self.cli.output.get_size()

        margin = self.get_reserved_space() + self.get_prompt(
            self.prompt_format).count('\n') + 1
        if special.is_timing_enabled():
            margin += 1
        if status:
            margin += 1 + status.count('\n')

        for i, line in enumerate(output.splitlines(), 1):
            if len(line) > size.columns or i > (size.rows - margin):
                return False

        return True

    def output(self, output, status=None):
        """Output text to stdout or a pager command.

        The status text is not outputted to pager or files.

        The message will be logged in the audit log, if enabled. The
        message will be written to the tee file, if enabled. The
        message will be written to the output file, if enabled.

        """
        if output:
            self.log_output(output)
            special.write_tee(output)
            special.write_once(output)

            if (self.explicit_pager
                    or (special.is_pager_enabled()
                        and not self.output_fits_on_screen(output, status))):
                click.echo_via_pager(output)
            else:
                click.secho(output)

        if status:
            self.log_output(status)
            click.secho(status)

    def configure_pager(self):
        # Provide sane defaults for less if they are empty.
        if not os.environ.get('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'])
            self.explicit_pager = True
        else:
            self.explicit_pager = False

        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, {
                'smart_completion': self.smart_completion,
                'supported_formats': self.formatter.supported_formats
            })

        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
        host = self.login_path if self.login_path and self.login_path_as_host else sqlexecute.host
        now = datetime.now()
        string = string.replace('\\u', sqlexecute.user or '(none)')
        string = string.replace('\\h', 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")
        string = string.replace('\\D', now.strftime('%a %b %d %H:%M:%S %Y'))
        string = string.replace('\\m', now.strftime('%M'))
        string = string.replace('\\P', now.strftime('%p'))
        string = string.replace('\\R', now.strftime('%H'))
        string = string.replace('\\r', now.strftime('%I'))
        string = string.replace('\\s', now.strftime('%S'))
        string = string.replace('\\p', str(sqlexecute.port))
        string = string.replace('\\_', ' ')
        return string

    def run_query(self, query, new_line=True):
        """Runs *query*."""
        results = self.sqlexecute.run(query)
        for result in results:
            title, cur, headers, status = result
            output = self.format_output(title, cur, headers)
            for line in output:
                click.echo(line, nl=new_line)

    def format_output(self,
                      title,
                      cur,
                      headers,
                      expanded=False,
                      max_width=None):
        expanded = expanded or self.formatter.format_name == 'vertical'
        output = []

        if title:  # Only print the title if it's not None.
            output.append(title)

        if cur:
            rows = list(cur)
            formatted = self.formatter.format_output(
                rows, headers, format_name='vertical' if expanded else None)

            if (not expanded and max_width and rows
                    and content_exceeds_width(rows[0], max_width) and headers):
                formatted = self.formatter.format_output(
                    rows, headers, format_name='vertical')

            output.append(formatted)

        return output

    def get_reserved_space(self):
        """Get the number of lines to reserve for the completion menu."""
        reserved_space_ratio = .45
        max_reserved_space = 8
        _, height = click.get_terminal_size()
        return min(round(height * reserved_space_ratio), max_reserved_space)

    def get_last_query(self):
        """Get the last query executed or None."""
        return self.query_history[-1][0] if self.query_history else None
Esempio n. 24
0
File: main.py Progetto: pranav/mycli
    def run_cli(self):
        sqlexecute = self.sqlexecute
        logger = self.logger
        original_less_opts = self.adjust_less_opts()
        self.set_pager_from_config()

        self.initialize_completions()
        completer = self.completer

        def set_key_bindings(value):
            if value not in ('emacs', 'vi'):
                value = 'emacs'
            self.key_bindings = value

        project_root = os.path.dirname(PACKAGE_ROOT)
        author_file = os.path.join(project_root, 'AUTHORS')
        sponsor_file = os.path.join(project_root, 'SPONSORS')

        key_binding_manager = mycli_bindings(get_key_bindings=lambda: self.key_bindings,
                                             set_key_bindings=set_key_bindings)
        print('Version:', __version__)
        print('Chat: https://gitter.im/dbcli/mycli')
        print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
        print('Home: http://mycli.net')
        print('Thanks to the contributor -', thanks_picker([author_file, sponsor_file]))

        def prompt_tokens(cli):
            return [(Token.Prompt, self.get_prompt(self.prompt_format))]

        get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.key_bindings)
        layout = create_default_layout(lexer=MyCliLexer,
                                       reserve_space_for_menu=True,
                                       multiline=True,
                                       get_prompt_tokens=prompt_tokens,
                                       get_bottom_toolbar_tokens=get_toolbar_tokens,
                                       display_completions_in_columns=self.wider_completion_menu,
                                       extra_input_processors=[
                                           ConditionalProcessor(
                                               processor=HighlightMatchingBracketProcessor(chars='[](){}'),
                                               filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
                                       ])
        buf = CLIBuffer(always_multiline=self.multi_line, completer=completer,
                history=FileHistory(os.path.expanduser('~/.mycli-history')),
                complete_while_typing=Always())

        application = Application(style=style_factory(self.syntax_style, self.cli_style),
                                  layout=layout, buffer=buf,
                                  key_bindings_registry=key_binding_manager.registry,
                                  on_exit=AbortAction.RAISE_EXCEPTION,
                                  ignore_case=True)
        cli = CommandLineInterface(application=application, eventloop=create_eventloop())

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

                special.set_expanded_output(False)

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

                try:
                    document = self.handle_editor_command(cli, document)
                except RuntimeError as e:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    self.output(str(e), err=True, fg='red')
                    continue
                if self.destructive_warning:
                    destroy = confirm_destructive_query(document.text)
                    if destroy is None:
                        pass  # Query was not destructive. Nothing to do here.
                    elif destroy is True:
                        self.output('Your call!')
                    else:
                        self.output('Wise choice!')
                        continue

                # Keep track of whether or not the query is mutating. In case
                # of a multi-statement query, the overall query is considered
                # mutating if any one of the component statements is mutating
                mutating = False

                try:
                    logger.debug('sql: %r', document.text)
                    if self.logfile:
                        self.logfile.write('\n# %s\n' % datetime.now())
                        self.logfile.write(document.text)
                        self.logfile.write('\n')
                    successful = False
                    start = time()
                    res = sqlexecute.run(document.text)
                    duration = time() - start
                    successful = True
                    output = []
                    total = 0
                    for title, cur, headers, status in res:
                        logger.debug("headers: %r", headers)
                        logger.debug("rows: %r", cur)
                        logger.debug("status: %r", status)
                        start = time()
                        threshold = 1000
                        if (is_select(status) and
                                cur and cur.rowcount > threshold):
                            self.output('The result set has more than %s rows.'
                                    % threshold, fg='red')
                            if not click.confirm('Do you want to continue?'):
                                self.output("Aborted!", err=True, fg='red')
                                break
                        output.extend(format_output(title, cur, headers,
                            status, self.table_format))
                        end = time()
                        total += end - start
                        mutating = mutating or is_mutating(status)
                except KeyboardInterrupt:
                    # Restart connection to the database
                    sqlexecute.connect()
                    logger.debug("cancelled query, sql: %r", document.text)
                    self.output("cancelled query", err=True, fg='red')
                except NotImplementedError:
                    self.output('Not Yet Implemented.', fg="yellow")
                except OperationalError as e:
                    logger.debug("Exception: %r", e)
                    reconnect = True
                    if (e.args[0] in (2003, 2006, 2013)):
                        reconnect = click.prompt('Connection reset. Reconnect (Y/n)',
                                show_default=False, type=bool, default=True)
                        if reconnect:
                            logger.debug('Attempting to reconnect.')
                            try:
                                sqlexecute.connect()
                                logger.debug('Reconnected successfully.')
                                self.output('Reconnected!\nTry the command again.', fg='green')
                            except OperationalError as e:
                                logger.debug('Reconnect failed. e: %r', e)
                                self.output(str(e), err=True, fg='red')
                                continue  # If reconnection failed, don't proceed further.
                        else:  # If user chooses not to reconnect, don't proceed further.
                            continue
                    else:
                        logger.error("sql: %r, error: %r", document.text, e)
                        logger.error("traceback: %r", traceback.format_exc())
                        self.output(str(e), err=True, fg='red')
                except Exception as e:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    self.output(str(e), err=True, fg='red')
                else:
                    try:
                        self.output_via_pager('\n'.join(output))
                    except KeyboardInterrupt:
                        pass
                    if special.is_timing_enabled():
                        self.output('Command Time: %0.03fs' % duration)
                        self.output('Format Time: %0.03fs' % total)

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_dynamic_completions()

                query = Query(document.text, successful, mutating)
                self.query_history.append(query)

        except EOFError:
            self.output('Goodbye!')
        finally:  # Reset the less opts back to original.
            logger.debug('Restoring env var LESS to %r.', original_less_opts)
            os.environ['LESS'] = original_less_opts
            os.environ['PAGER'] = special.get_original_pager()
Esempio n. 25
0
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!')
Esempio n. 26
0
    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.")
Esempio n. 27
0
        Token.Placeholder: "#888888",
        Token.Placeholder.Variable: "#888888",
        Token.Placeholder.Bracket: "bold #ff7777",
        Token.Placeholder.Separator: "#ee7777",
        Token.Aborted: '#aaaaaa',
        Token.Menu.Completions.Completion.Current: 'bg:#00aaaa #000000',
        Token.Menu.Completions.Completion: 'bg:#008888 #ffffff',
        Token.Menu.Completions.ProgressButton: 'bg:#003333',
        Token.Menu.Completions.ProgressBar: 'bg:#00aaaa',
    }


if __name__ == '__main__':
    cli = CommandLineInterface(layout=Layout(
        before_input=DefaultPrompt('Calculate: '),
        after_input=CompletionHint(grammar),
        menus=[CompletionsMenu()]),
                               line=Line(completer=ShellCompleter(grammar)),
                               style=ExampleStyle)

    try:
        # REPL loop.
        while True:
            # Read input and parse the result.
            try:
                document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
                parse_info = get_parse_info(grammar, document)
            except InvalidCommandException:
                print('Invalid command\n')
                continue

            # For debugging only: show the data structure that we got.
Esempio n. 28
0
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.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:
            for key, value in self.settings.items():
                self.client.query('SET {}={}'.format(key, value), fmt='Null')

            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')
        # forcefully disable `highlight_output` in (u)rxvt (https://github.com/hatarist/clickhouse-cli/issues/20)
        self.highlight_output = False if os.environ.get(
            'TERM', '').startswith('rxvt') else 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 os.environ.get(
            'CLICKHOUSE_HOST', '') or self.config.get('defaults',
                                                      'host') or '127.0.0.1'
        self.port = self.port or int(
            os.environ.get('CLICKHOUSE_PORT', '0')) or self.config.get(
                'defaults', 'port') or 8123
        self.user = self.user or os.environ.get(
            'CLICKHOUSE_USER', '') or self.config.get('defaults',
                                                      'user') or 'default'
        self.password = self.password or os.environ.get(
            'CLICKHOUSE_PASSWORD', '') or self.config.get(
                'defaults', 'password')
        self.database = self.database or os.environ.get(
            'CLICKHOUSE_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_output 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') and not self.echo.verbose:
            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):
        if not self.echo.verbose:
            return (0, 0)

        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()
Esempio n. 29
0
    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!')
Esempio n. 30
0
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)
Esempio n. 31
0
    def _build_cli(self, history):
        key_binding_manager = cli_bindings()

        def prompt_tokens(cli):
            prompt = self.get_prompt(self.prompt)
            if len(prompt) > self.MAX_LEN_PROMPT:
                prompt = self.get_prompt('\\r:\\d> ')
            return [(Token.Prompt, prompt)]

        def get_continuation_tokens(cli, width):
            prompt = self.get_prompt(self.prompt_continuation_format)
            token = (Token.Continuation, ' ' * (width - len(prompt)) + prompt)
            return [token]

        def show_suggestion_tip():
            return self.iterations < 2

        get_toolbar_tokens = create_toolbar_tokens_func(
            self.completion_refresher.is_refreshing, show_suggestion_tip)

        layout = create_prompt_layout(
            lexer=Lexer,
            multiline=True,
            get_prompt_tokens=prompt_tokens,
            get_continuation_tokens=get_continuation_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=False,
            extra_input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
            ],
            reserve_space_for_menu=self.get_reserved_space())

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

            if self.key_bindings == 'vi':
                editing_mode = EditingMode.VI
            else:
                editing_mode = EditingMode.EMACS

            application = Application(
                style=style_from_pygments(style_cls=self.output_style),
                layout=layout,
                buffer=buf,
                key_bindings_registry=key_binding_manager.registry,
                on_exit=AbortAction.RAISE_EXCEPTION,
                on_abort=AbortAction.RETRY,
                editing_mode=editing_mode,
                ignore_case=True)

            cli = CommandLineInterface(application=application,
                                       eventloop=create_eventloop())

            return cli
Esempio n. 32
0
    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)
Esempio n. 33
0
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_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,
                            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()
            self._process_command(document.text)
Esempio n. 34
0
    def run_cli(self):
        pgexecute = self.pgexecute
        prompt = '%s> ' % pgexecute.dbname
        logger = self.logger
        original_less_opts = self.adjust_less_opts()

        completer = self.completer
        self.refresh_completions()
        key_binding_manager = pgcli_bindings(self.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')

        layout = Layout(before_input=DefaultPrompt(prompt),
            menus=[CompletionsMenu(max_height=10)],
            lexer=PostgresLexer,
            bottom_toolbars=[PGToolbar(key_binding_manager)])
        buf = PGBuffer(always_multiline=self.multi_line, completer=completer,
                history=FileHistory(os.path.expanduser('~/.pgcli-history')))
        cli = CommandLineInterface(style=style_factory(self.syntax_style),
                layout=layout, buffer=buf,
                key_bindings_registry=key_binding_manager.registry)

        try:
            while True:
                cli.layout.before_input = DefaultPrompt(prompt)
                document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)

                # 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 Exit

                # 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)
                    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
                        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
                    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 pgspecial.TIMING_ENABLED:
                        print('Command Time:', duration)
                        print('Format Time:', total)
                finally:
                    for cur, _, _ in res:
                        if hasattr(cur, 'close'):
                            cur.close()

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    prompt = '%s> ' % pgexecute.dbname
                    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 Exit:
            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
Esempio n. 35
0
File: main.py Progetto: Hsuing/mycli
    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!')
Esempio n. 36
0
def main():
    manager = KeyBindingManager(enable_system_prompt=True)

    D = LayoutDimension
    layout = HSplit([
        VSplit([
            Window(width=D(min=15, max=30, preferred=30),
                   content=FillControl('a', token=Token.A)),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            Window(content=TokenListControl.static([(Token.HelloWorld,
                                                     lipsum)])),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            Window(content=BufferControl(lexer=PythonLexer,
                                         show_line_numbers=Always(),
                                         input_processors=[
                                             DefaultPrompt('python> '),
                                             AfterInput.static(
                                                 ' <python',
                                                 token=Token.AfterInput),
                                         ]), ),
            Window(width=D.exact(1),
                   content=FillControl('|', token=Token.Line)),
            HSplit([
                Window(width=D(max=40),
                       height=D.exact(4),
                       content=FillControl('b', token=Token.B)),
                Window(width=D(max=40),
                       content=FillControl('f', token=Token.F)),
                Window(width=D.exact(30),
                       height=D.exact(2),
                       content=FillControl('c', token=Token.C)),
            ]),
            #CompletionsMenu(),
        ]),
        Window(height=D.exact(1), content=FillControl('-', token=Token.Line)),
        Window(height=D.exact(3), content=FillControl('d', token=Token.D)),
        SystemToolbar(),
        ArgToolbar(),
        CompletionsToolbar(),
        SearchToolbar(),
    ])

    layout = FloatContainer(content=layout,
                            floats=[
                                Float(xcursor=True,
                                      ycursor=True,
                                      content=VSplit([
                                          Window(width=D.exact(5),
                                                 content=FillControl(
                                                     'f', token=Token.F)),
                                          CompletionsMenu(),
                                      ])),
                            ])

    eventloop = create_eventloop()
    cli = CommandLineInterface(eventloop=eventloop,
                               layout=layout,
                               style=TestStyle,
                               key_bindings_registry=manager.registry,
                               buffer=Buffer(is_multiline=Always(),
                                             completer=TestCompleter()))
    cli.read_input()
    eventloop.close()
Esempio n. 37
0
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:
                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()
            self.handle_exit(document)
            self.run_command(document)
Esempio n. 38
0
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
Esempio n. 39
0
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!')
Esempio n. 40
0
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
Esempio n. 41
0
    def run_cli(self):
        sqlexecute = self.sqlexecute
        logger = self.logger
        original_less_opts = self.adjust_less_opts()
        self.set_pager_from_config()

        self.initialize_completions()
        completer = self.completer

        def set_key_bindings(value):
            if value not in ("emacs", "vi"):
                value = "emacs"
            self.key_bindings = value

        project_root = os.path.dirname(PACKAGE_ROOT)
        author_file = os.path.join(project_root, "AUTHORS")
        sponsor_file = os.path.join(project_root, "SPONSORS")

        key_binding_manager = mycli_bindings(
            get_key_bindings=lambda: self.key_bindings, set_key_bindings=set_key_bindings
        )
        print("Version:", __version__)
        print("Chat: https://gitter.im/dbcli/mycli")
        print("Mail: https://groups.google.com/forum/#!forum/mycli-users")
        print("Home: http://mycli.net")
        print("Thanks to the contributor -", thanks_picker([author_file, sponsor_file]))

        def prompt_tokens(cli):
            return [(Token.Prompt, self.get_prompt(self.prompt_format))]

        get_toolbar_tokens = create_toolbar_tokens_func(lambda: self.key_bindings)
        layout = create_default_layout(
            lexer=MyCliLexer,
            reserve_space_for_menu=True,
            get_prompt_tokens=prompt_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            extra_input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(chars="[](){}"),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone(),
                )
            ],
        )
        buf = CLIBuffer(
            always_multiline=self.multi_line,
            completer=completer,
            history=FileHistory(os.path.expanduser("~/.mycli-history")),
            complete_while_typing=Always(),
        )

        application = Application(
            style=style_factory(self.syntax_style),
            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()
Esempio n. 42
0
    def run_cli(self):
        sqlexecute = self.sqlexecute
        logger = self.logger
        self.configure_pager()

        if self.smart_completion:
            self.refresh_completions()

        author_file = os.path.join(PACKAGE_ROOT, 'AUTHORS')
        sponsor_file = os.path.join(PACKAGE_ROOT, 'SPONSORS')

        key_binding_manager = mycli_bindings()

        if not self.less_chatty:
            print('Version:', __version__)
            print('Chat: https://gitter.im/dbcli/mycli')
            print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
            print('Home: http://mycli.net')
            print('Thanks to the contributor -',
                  thanks_picker([author_file, sponsor_file]))

        def prompt_tokens(cli):
            prompt = self.get_prompt(self.prompt_format)
            if self.prompt_format == self.default_prompt and len(
                    prompt) > self.max_len_prompt:
                prompt = self.get_prompt('\\d> ')
            return [(Token.Prompt, prompt)]

        def get_continuation_tokens(cli, width):
            continuation_prompt = self.get_prompt(
                self.prompt_continuation_format)
            return [(Token.Continuation, ' ' *
                     (width - len(continuation_prompt)) + continuation_prompt)]

        def one_iteration(document=None):
            if document is None:
                document = self.cli.run()

                special.set_expanded_output(False)

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

            if not document.text.strip():
                return

            if self.destructive_warning:
                destroy = confirm_destructive_query(document.text)
                if destroy is None:
                    pass  # Query was not destructive. Nothing to do here.
                elif destroy is True:
                    self.echo('Your call!')
                else:
                    self.echo('Wise choice!')
                    return

            # Keep track of whether or not the query is mutating. In case
            # of a multi-statement query, the overall query is considered
            # mutating if any one of the component statements is mutating
            mutating = False

            try:
                logger.debug('sql: %r', document.text)

                special.write_tee(
                    self.get_prompt(self.prompt_format) + document.text)
                if self.logfile:
                    self.logfile.write('\n# %s\n' % datetime.now())
                    self.logfile.write(document.text)
                    self.logfile.write('\n')

                successful = False
                start = time()
                res = sqlexecute.run(document.text)
                successful = True
                result_count = 0
                for title, cur, headers, status in res:
                    logger.debug("headers: %r", headers)
                    logger.debug("rows: %r", cur)
                    logger.debug("status: %r", status)
                    threshold = 1000
                    if (is_select(status) and cur
                            and cur.rowcount > threshold):
                        self.echo(
                            'The result set has more than {} rows.'.format(
                                threshold),
                            fg='red')
                        if not click.confirm('Do you want to continue?'):
                            self.echo("Aborted!", err=True, fg='red')
                            break

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

                    formatted = self.format_output(
                        title, cur, headers, special.is_expanded_output(),
                        max_width)

                    t = time() - start
                    try:
                        if result_count > 0:
                            self.echo('')
                        try:
                            self.output('\n'.join(formatted), status)
                        except KeyboardInterrupt:
                            pass
                        if special.is_timing_enabled():
                            self.echo('Time: %0.03fs' % t)
                    except KeyboardInterrupt:
                        pass

                    start = time()
                    result_count += 1
                    mutating = mutating or is_mutating(status)
                special.unset_once_if_written()
            except EOFError as e:
                raise e
            except KeyboardInterrupt:
                # get last connection id
                connection_id_to_kill = sqlexecute.connection_id
                logger.debug("connection id to kill: %r",
                             connection_id_to_kill)
                # Restart connection to the database
                sqlexecute.connect()
                try:
                    for title, cur, headers, status in sqlexecute.run(
                            'kill %s' % connection_id_to_kill):
                        status_str = str(status).lower()
                        if status_str.find('ok') > -1:
                            logger.debug(
                                "cancelled query, connection id: %r, sql: %r",
                                connection_id_to_kill, document.text)
                            self.echo("cancelled query", err=True, fg='red')
                except Exception as e:
                    self.echo(
                        'Encountered error while cancelling query: {}'.format(
                            e),
                        err=True,
                        fg='red')
            except NotImplementedError:
                self.echo('Not Yet Implemented.', fg="yellow")
            except OperationalError as e:
                logger.debug("Exception: %r", e)
                if (e.args[0] in (2003, 2006, 2013)):
                    logger.debug('Attempting to reconnect.')
                    self.echo('Reconnecting...', fg='yellow')
                    try:
                        sqlexecute.connect()
                        logger.debug('Reconnected successfully.')
                        one_iteration(document)
                        return  # OK to just return, cuz the recursion call runs to the end.
                    except OperationalError as e:
                        logger.debug('Reconnect failed. e: %r', e)
                        self.echo(str(e), err=True, fg='red')
                        # If reconnection failed, don't proceed further.
                        return
                else:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    self.echo(str(e), err=True, fg='red')
            except Exception as e:
                logger.error("sql: %r, error: %r", document.text, e)
                logger.error("traceback: %r", traceback.format_exc())
                self.echo(str(e), err=True, fg='red')
            else:
                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_completions(
                        reset=need_completion_reset(document.text))
            finally:
                if self.logfile is False:
                    self.echo("Warning: This query was not logged.",
                              err=True,
                              fg='red')
            query = Query(document.text, successful, mutating)
            self.query_history.append(query)

        get_toolbar_tokens = create_toolbar_tokens_func(
            self.completion_refresher.is_refreshing)

        layout = create_prompt_layout(
            lexer=MyCliLexer,
            multiline=True,
            get_prompt_tokens=prompt_tokens,
            get_continuation_tokens=get_continuation_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            extra_input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone())
            ],
            reserve_space_for_menu=self.get_reserved_space())
        with self._completer_lock:
            buf = CLIBuffer(always_multiline=self.multi_line,
                            completer=self.completer,
                            history=FileHistory(
                                os.path.expanduser(
                                    os.environ.get('MYCLI_HISTFILE',
                                                   '~/.mycli-history'))),
                            complete_while_typing=Always(),
                            accept_action=AcceptAction.RETURN_DOCUMENT)

            if self.key_bindings == 'vi':
                editing_mode = EditingMode.VI
            else:
                editing_mode = EditingMode.EMACS

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

        try:
            while True:
                one_iteration()
        except EOFError:
            special.close_tee()
            if not self.less_chatty:
                self.echo('Goodbye!')
Esempio n. 43
0
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
Esempio n. 44
0
    def run_cli(self):
        pgexecute = self.pgexecute
        logger = self.logger
        original_less_opts = self.adjust_less_opts()

        self.refresh_completions()

        def set_vi_mode(value):
            self.vi_mode = value

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

        print('Version:', __version__)
        print('Chat: https://gitter.im/dbcli/pgcli')
        print('Mail: https://groups.google.com/forum/#!forum/pgcli')
        print('Home: http://pgcli.com')

        def prompt_tokens(cli):
            return [(Token.Prompt, '%s> ' % pgexecute.dbname)]

        get_toolbar_tokens = create_toolbar_tokens_func(
            lambda: self.vi_mode,
            lambda: self.completion_refresher.is_refreshing())

        layout = create_default_layout(
            lexer=PostgresLexer,
            reserve_space_for_menu=True,
            get_prompt_tokens=prompt_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            multiline=True,
            extra_input_processors=[
                # Highlight matching brackets while editing.
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
            ])
        history_file = self.config['main']['history_file']
        with self._completer_lock:
            buf = PGBuffer(always_multiline=self.multi_line,
                           completer=self.completer,
                           history=FileHistory(
                               os.path.expanduser(history_file)),
                           complete_while_typing=Always())

            application = Application(
                style=style_factory(self.syntax_style, self.cli_style),
                layout=layout,
                buffer=buf,
                key_bindings_registry=key_binding_manager.registry,
                on_exit=AbortAction.RAISE_EXCEPTION,
                ignore_case=True)
            self.cli = CommandLineInterface(application=application,
                                            eventloop=create_eventloop())

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

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

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

                # Keep track of whether or not the query is mutating. In case
                # of a multi-statement query, the overall query is considered
                # mutating if any one of the component statements is mutating
                mutating = False

                try:
                    logger.debug('sql: %r', document.text)
                    successful = False
                    # Initialized to [] because res might never get initialized
                    # if an exception occurs in pgexecute.run(). Which causes
                    # finally clause to fail.
                    res = []
                    # Run the query.
                    start = time()
                    res = pgexecute.run(document.text, self.pgspecial)
                    successful = True
                    output = []
                    total = 0
                    for title, cur, headers, status in res:
                        logger.debug("headers: %r", headers)
                        logger.debug("rows: %r", cur)
                        logger.debug("status: %r", status)
                        threshold = 1000
                        if (is_select(status) and cur
                                and cur.rowcount > threshold):
                            click.secho(
                                'The result set has more than %s rows.' %
                                threshold,
                                fg='red')
                            if not click.confirm('Do you want to continue?'):
                                click.secho("Aborted!", err=True, fg='red')
                                break

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

                        formatted = format_output(
                            title, cur, headers, status, self.table_format,
                            self.pgspecial.expanded_output, max_width)
                        output.extend(formatted)
                        end = time()
                        total += end - start
                        mutating = mutating or is_mutating(status)

                except KeyboardInterrupt:
                    # Restart connection to the database
                    pgexecute.connect()
                    logger.debug("cancelled query, sql: %r", document.text)
                    click.secho("cancelled query", err=True, fg='red')
                except NotImplementedError:
                    click.secho('Not Yet Implemented.', fg="yellow")
                except OperationalError as e:
                    reconnect = True
                    if ('server closed the connection'
                            in utf8tounicode(e.args[0])):
                        reconnect = click.prompt(
                            'Connection reset. Reconnect (Y/n)',
                            show_default=False,
                            type=bool,
                            default=True)
                        if reconnect:
                            try:
                                pgexecute.connect()
                                click.secho(
                                    'Reconnected!\nTry the command again.',
                                    fg='green')
                            except OperationalError as e:
                                click.secho(str(e), err=True, fg='red')
                    else:
                        logger.error("sql: %r, error: %r", document.text, e)
                        logger.error("traceback: %r", traceback.format_exc())
                        click.secho(str(e), err=True, fg='red')
                except Exception as e:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    click.secho(str(e), err=True, fg='red')
                else:
                    try:
                        click.echo_via_pager('\n'.join(output))
                    except KeyboardInterrupt:
                        pass
                    if self.pgspecial.timing_enabled:
                        print('Time: %0.03fs' % total)

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_completions()

                # Refresh search_path to set default schema.
                if need_search_path_refresh(document.text):
                    logger.debug('Refreshing search path')
                    with self._completer_lock:
                        self.completer.set_search_path(pgexecute.search_path())
                    logger.debug('Search path: %r', self.completer.search_path)

                query = Query(document.text, successful, mutating)
                self.query_history.append(query)

        except EOFError:
            print('Goodbye!')
        finally:  # Reset the less opts back to original.
            logger.debug('Restoring env var LESS to %r.', original_less_opts)
            os.environ['LESS'] = original_less_opts
Esempio n. 45
0
File: main.py Progetto: nosun/pgcli
    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
class StatusToolbar(Toolbar):
    """
    Status toolbar showing the current location.
    """
    def __init__(self):
        super(StatusToolbar, self).__init__(Token.StatusBar)

    def get_tokens(self, cli, width):
        return [(self.token, ' '), (self.token, os.getcwd())]


if __name__ == '__main__':
    cli = CommandLineInterface(
        layout=Layout(
            before_input=DefaultPrompt('Shell> '),
            #after_input=CompletionHint(grammar),
            menus=[CompletionsMenu()],
            bottom_toolbars=[StatusToolbar()]),
        line=Line(completer=ShellCompleter(grammar)),
        style=ExampleStyle)

    try:
        while True:
            try:
                document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
                parse_info = get_parse_info(grammar, document)
            except InvalidCommandException:
                os.system(document.text)
                continue
            else:
                params = parse_info.get_variables()
Esempio n. 47
0
    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!")
Esempio n. 48
0
    def run_cli(self):
        sqlexecute = self.sqlexecute
        logger = self.logger
        self.configure_pager()

        self.refresh_completions()

        project_root = os.path.dirname(PACKAGE_ROOT)
        author_file = os.path.join(project_root, 'AUTHORS')
        sponsor_file = os.path.join(project_root, 'SPONSORS')

        key_binding_manager = mycli_bindings()

        if not self.less_chatty:
            print('Version:', __version__)
            print('Chat: https://gitter.im/dbcli/mycli')
            print('Mail: https://groups.google.com/forum/#!forum/mycli-users')
            print('Home: http://mycli.net')
            print('Thanks to the contributor -',
                  thanks_picker([author_file, sponsor_file]))

        def prompt_tokens(cli):
            return [(Token.Prompt, self.get_prompt(self.prompt_format))]

        def get_continuation_tokens(cli, width):
            continuation_prompt = self.get_prompt(
                self.prompt_continuation_format)
            return [(Token.Continuation, ' ' *
                     (width - len(continuation_prompt)) + continuation_prompt)]

        def one_iteration(document=None):
            if document is None:
                document = self.cli.run(reset_current_buffer=True)

                special.set_expanded_output(False)

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

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

            if self.destructive_warning:
                destroy = confirm_destructive_query(document.text)
                if destroy is None:
                    pass  # Query was not destructive. Nothing to do here.
                elif destroy is True:
                    self.output('Your call!')
                else:
                    self.output('Wise choice!')
                    return

            # Keep track of whether or not the query is mutating. In case
            # of a multi-statement query, the overall query is considered
            # mutating if any one of the component statements is mutating
            mutating = False

            try:
                logger.debug('sql: %r', document.text)

                if self.logfile:
                    self.logfile.write('\n# %s\n' % datetime.now())
                    self.logfile.write(document.text)
                    self.logfile.write('\n')

                successful = False
                start = time()
                res = sqlexecute.run(document.text)
                successful = True
                output = []
                total = 0
                for title, cur, headers, status in res:
                    logger.debug("headers: %r", headers)
                    logger.debug("rows: %r", cur)
                    logger.debug("status: %r", status)
                    threshold = 1000
                    if (is_select(status) and cur
                            and cur.rowcount > threshold):
                        self.output('The result set has more than %s rows.' %
                                    threshold,
                                    fg='red')
                        if not click.confirm('Do you want to continue?'):
                            self.output("Aborted!", err=True, fg='red')
                            break

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

                    formatted = format_output(title, cur, headers, status,
                                              self.table_format,
                                              special.is_expanded_output(),
                                              max_width)

                    output.extend(formatted)
                    end = time()
                    total += end - start
                    mutating = mutating or is_mutating(status)
            except UnicodeDecodeError as e:
                import pymysql
                if pymysql.VERSION < (0, 6, 7):
                    message = (
                        'You are running an older version of pymysql.\n'
                        'Please upgrade to 0.6.7 or above to view binary data.\n'
                        'Try \'pip install -U pymysql\'.')
                    self.output(message)
                else:
                    raise e
            except KeyboardInterrupt:
                # Restart connection to the database
                sqlexecute.connect()
                logger.debug("cancelled query, sql: %r", document.text)
                self.output("cancelled query", err=True, fg='red')
            except NotImplementedError:
                self.output('Not Yet Implemented.', fg="yellow")
            except OperationalError as e:
                logger.debug("Exception: %r", e)
                if (e.args[0] in (2003, 2006, 2013)):
                    logger.debug('Attempting to reconnect.')
                    self.output('Reconnecting...', fg='yellow')
                    try:
                        sqlexecute.connect()
                        logger.debug('Reconnected successfully.')
                        one_iteration(document)
                        return  # OK to just return, cuz the recursion call runs to the end.
                    except OperationalError as e:
                        logger.debug('Reconnect failed. e: %r', e)
                        self.output(str(e), err=True, fg='red')
                        return  # If reconnection failed, don't proceed further.
                else:
                    logger.error("sql: %r, error: %r", document.text, e)
                    logger.error("traceback: %r", traceback.format_exc())
                    self.output(str(e), err=True, fg='red')
            except Exception as e:
                logger.error("sql: %r, error: %r", document.text, e)
                logger.error("traceback: %r", traceback.format_exc())
                self.output(str(e), err=True, fg='red')
            else:
                try:
                    if special.is_pager_enabled():
                        self.output_via_pager('\n'.join(output))
                    else:
                        self.output('\n'.join(output))
                except KeyboardInterrupt:
                    pass
                if special.is_timing_enabled():
                    self.output('Time: %0.03fs' % total)

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    self.refresh_completions(
                        reset=need_completion_reset(document.text))
            finally:
                if self.logfile is False:
                    self.output("Warning: This query was not logged.",
                                err=True,
                                fg='red')
            query = Query(document.text, successful, mutating)
            self.query_history.append(query)

        get_toolbar_tokens = create_toolbar_tokens_func(
            self.completion_refresher.is_refreshing)

        layout = create_prompt_layout(
            lexer=MyCliLexer,
            multiline=True,
            get_prompt_tokens=prompt_tokens,
            get_continuation_tokens=get_continuation_tokens,
            get_bottom_toolbar_tokens=get_toolbar_tokens,
            display_completions_in_columns=self.wider_completion_menu,
            extra_input_processors=[
                ConditionalProcessor(
                    processor=HighlightMatchingBracketProcessor(
                        chars='[](){}'),
                    filter=HasFocus(DEFAULT_BUFFER) & ~IsDone()),
            ])
        with self._completer_lock:
            buf = CLIBuffer(always_multiline=self.multi_line,
                            completer=self.completer,
                            history=FileHistory(
                                os.path.expanduser(
                                    os.environ.get('MYCLI_HISTFILE',
                                                   '~/.mycli-history'))),
                            complete_while_typing=Always(),
                            accept_action=AcceptAction.RETURN_DOCUMENT)

            if self.key_bindings == 'vi':
                editing_mode = EditingMode.VI
            else:
                editing_mode = EditingMode.EMACS

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

        try:
            while True:
                one_iteration()
        except EOFError:
            if not self.less_chatty:
                self.output('Goodbye!')
def main():
    cli = CommandLineInterface(layout=layout, style=TestStyle, line=Line(validator=EmailValidator()))

    document = cli.read_input()
    print('You said: ' + document.text)
Esempio n. 50
0
    def _build_cli(self, history):
        def set_vi_mode(value):
            self.vi_mode = value

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

        def prompt_tokens(_):
            if self.dsn_alias and self.prompt_dsn_format is not None:
                prompt_format = self.prompt_dsn_format
            else:
                prompt_format = self.prompt_format

            prompt = self.get_prompt(prompt_format)

            if (prompt_format == self.default_prompt
                    and len(prompt) > self.max_len_prompt):
                prompt = self.get_prompt('\\d> ')

            return [(Token.Prompt, prompt)]

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

        get_toolbar_tokens = create_toolbar_tokens_func(
            lambda: self.vi_mode, self.completion_refresher.is_refreshing,
            self.pgexecute.failed_transaction,
            self.pgexecute.valid_transaction)

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

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

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

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

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

            return cli
Esempio n. 51
0
File: main.py Progetto: Hsuing/mycli
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
Esempio n. 52
0
    def run_cli(self):
        pgexecute = self.pgexecute
        prompt = '%s> ' % pgexecute.dbname
        logger = self.logger
        original_less_opts = self.adjust_less_opts()

        completer = self.completer
        self.refresh_completions()
        key_binding_manager = pgcli_bindings(self.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')

        layout = Layout(before_input=DefaultPrompt(prompt),
                        menus=[CompletionsMenu(max_height=10)],
                        lexer=PostgresLexer,
                        bottom_toolbars=[PGToolbar(key_binding_manager)])
        buf = PGBuffer(always_multiline=self.multi_line,
                       completer=completer,
                       history=FileHistory(
                           os.path.expanduser('~/.pgcli-history')))
        cli = CommandLineInterface(
            style=style_factory(self.syntax_style),
            layout=layout,
            buffer=buf,
            key_bindings_registry=key_binding_manager.registry)

        try:
            while True:
                cli.layout.before_input = DefaultPrompt(prompt)
                document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)

                # 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 Exit

                # 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)
                    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
                        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
                    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 pgspecial.TIMING_ENABLED:
                        print('Command Time:', duration)
                        print('Format Time:', total)
                finally:
                    for cur, _, _ in res:
                        if hasattr(cur, 'close'):
                            cur.close()

                # Refresh the table names and column names if necessary.
                if need_completion_refresh(document.text):
                    prompt = '%s> ' % pgexecute.dbname
                    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 Exit:
            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