2014-10-12 17:31:54 +00:00
|
|
|
#!/usr/bin/env python
|
|
|
|
from __future__ import unicode_literals
|
2014-11-21 07:15:50 +00:00
|
|
|
from __future__ import print_function
|
2014-10-12 22:07:34 +00:00
|
|
|
|
2014-11-27 23:02:54 +00:00
|
|
|
import os.path
|
2014-11-23 23:31:34 +00:00
|
|
|
|
2014-10-12 22:07:34 +00:00
|
|
|
import click
|
2014-10-12 17:31:54 +00:00
|
|
|
|
2014-10-12 17:45:35 +00:00
|
|
|
from prompt_toolkit import CommandLineInterface, AbortAction, Exit
|
2014-10-12 17:31:54 +00:00
|
|
|
from prompt_toolkit.layout import Layout
|
|
|
|
from prompt_toolkit.layout.prompt import DefaultPrompt
|
2014-11-21 07:15:50 +00:00
|
|
|
from prompt_toolkit.layout.menus import CompletionsMenu
|
2014-11-23 23:31:34 +00:00
|
|
|
from prompt_toolkit.history import FileHistory
|
2014-11-24 07:30:17 +00:00
|
|
|
from pygments.lexers.sql import SqlLexer
|
2014-12-05 16:56:59 +00:00
|
|
|
import sqlparse
|
|
|
|
|
2014-12-08 08:43:21 +00:00
|
|
|
from .packages.tabulate import tabulate
|
|
|
|
from .packages.pgspecial import COMMANDS
|
2014-12-05 16:56:59 +00:00
|
|
|
from .pgcompleter import PGCompleter
|
|
|
|
from .pgstyle import PGStyle
|
|
|
|
from .pgexecute import PGExecute
|
|
|
|
from .pgline import PGLine
|
|
|
|
from .config import write_default_config, load_config
|
2014-10-12 17:31:54 +00:00
|
|
|
|
2014-10-12 22:07:34 +00:00
|
|
|
@click.command()
|
2014-11-21 07:15:50 +00:00
|
|
|
@click.option('-h', '--host', default='localhost')
|
2014-11-22 07:43:11 +00:00
|
|
|
@click.option('-p', '--port', default=5432)
|
2014-11-21 07:15:50 +00:00
|
|
|
@click.option('-U', '--user', prompt=True, envvar='USER')
|
2014-11-23 08:09:00 +00:00
|
|
|
@click.password_option('-W', '--password', default='',
|
2014-11-21 07:15:50 +00:00
|
|
|
confirmation_prompt=False)
|
|
|
|
@click.argument('database', envvar='USER')
|
2014-11-26 08:02:34 +00:00
|
|
|
def cli(database, user, password, host, port):
|
2014-11-27 23:02:54 +00:00
|
|
|
|
|
|
|
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.
|
2014-11-23 08:09:00 +00:00
|
|
|
try:
|
|
|
|
pgexecute = PGExecute(database, user, password, host, port)
|
|
|
|
except Exception as e:
|
|
|
|
click.secho(e.message, err=True, fg='red')
|
|
|
|
exit(1)
|
2014-12-08 06:30:24 +00:00
|
|
|
layout = Layout(before_input=DefaultPrompt('%s> ' % pgexecute.dbname),
|
2014-11-21 07:15:50 +00:00
|
|
|
menus=[CompletionsMenu()],
|
2014-10-12 17:45:35 +00:00
|
|
|
lexer=SqlLexer)
|
2014-11-28 06:06:31 +00:00
|
|
|
completer = PGCompleter(config.getboolean('main', 'smart_completion'))
|
2014-12-08 08:43:21 +00:00
|
|
|
completer.extend_special_commands(COMMANDS.keys())
|
2014-11-25 08:35:28 +00:00
|
|
|
completer.extend_table_names(pgexecute.tables())
|
|
|
|
completer.extend_column_names(pgexecute.all_columns())
|
2014-12-02 05:55:13 +00:00
|
|
|
line = PGLine(completer=completer,
|
2014-11-27 23:02:54 +00:00
|
|
|
history=FileHistory(os.path.expanduser('~/.pgcli-history')))
|
2014-11-21 07:15:50 +00:00
|
|
|
cli = CommandLineInterface(style=PGStyle, layout=layout, line=line)
|
2014-10-12 17:31:54 +00:00
|
|
|
|
2014-10-12 17:45:35 +00:00
|
|
|
try:
|
|
|
|
while True:
|
2014-12-08 06:30:24 +00:00
|
|
|
cli.layout.before_input = DefaultPrompt('%s> ' % pgexecute.dbname)
|
2014-10-12 17:45:35 +00:00
|
|
|
document = cli.read_input(on_exit=AbortAction.RAISE_EXCEPTION)
|
2014-12-06 06:04:39 +00:00
|
|
|
|
|
|
|
# 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.
|
2014-12-08 06:30:24 +00:00
|
|
|
if (document.text.strip().lower() == 'exit'
|
2014-12-08 08:43:21 +00:00
|
|
|
or document.text.strip().lower() == 'quit'
|
|
|
|
or document.text.strip() == '\q'):
|
2014-12-06 06:04:39 +00:00
|
|
|
raise Exit
|
2014-11-23 08:09:00 +00:00
|
|
|
try:
|
2014-12-09 06:59:25 +00:00
|
|
|
res = pgexecute.run(document.text)
|
|
|
|
for rows, headers, status in res:
|
|
|
|
if rows:
|
|
|
|
print(tabulate(rows, headers, tablefmt='psql'))
|
|
|
|
print(status)
|
2014-11-23 08:09:00 +00:00
|
|
|
except Exception as e:
|
|
|
|
click.secho(e.message, err=True, fg='red')
|
2014-12-05 16:56:59 +00:00
|
|
|
|
|
|
|
# Refresh the table names and column names if necessary.
|
|
|
|
if need_completion_refresh(document.text):
|
2014-12-08 08:43:21 +00:00
|
|
|
completer.reset_completions()
|
2014-12-05 16:56:59 +00:00
|
|
|
completer.extend_table_names(pgexecute.tables())
|
|
|
|
completer.extend_column_names(pgexecute.all_columns())
|
2014-10-12 17:45:35 +00:00
|
|
|
except Exit:
|
2014-11-21 07:15:50 +00:00
|
|
|
print ('GoodBye!')
|
2014-12-05 16:56:59 +00:00
|
|
|
|
|
|
|
def need_completion_refresh(sql):
|
2014-12-08 08:43:21 +00:00
|
|
|
first_token = sql.split()[0]
|
|
|
|
return first_token in ('alter', 'create', 'use', '\c', 'drop')
|