class CLI: def __init__(self, host, port, user, password, database, settings, format, format_stdin, multiline, stacktrace): self.config = None self.host = host self.port = port self.user = user self.password = password self.database = database self.settings = {k: v[0] for k, v in parse_qs(settings).items()} self.format = format self.format_stdin = format_stdin self.multiline = multiline self.stacktrace = stacktrace self.server_version = None self.query_ids = [] self.client = None self.echo = Echo(verbose=True, colors=True) self.progress = None self.metadata = {} def connect(self): self.scheme = 'http' if '://' in self.host: u = urlparse(self.host, allow_fragments=False) self.host = u.hostname self.port = u.port or self.port self.scheme = u.scheme self.url = '{scheme}://{host}:{port}/'.format(scheme=self.scheme, host=self.host, port=self.port) self.client = Client( self.url, self.user, self.password, self.database, self.settings, self.stacktrace, self.conn_timeout, self.conn_timeout_retry, self.conn_timeout_retry_delay, ) self.echo.print("Connecting to {host}:{port}".format(host=self.host, port=self.port)) try: response = self.client.query('SELECT version();', fmt='TabSeparated') except TimeoutError: self.echo.error("Error: Connection timeout.") return False except ConnectionError: self.echo.error("Error: Failed to connect.") return False except DBException as e: self.echo.error("Error:") self.echo.error(e.error) if self.stacktrace and e.stacktrace: self.echo.print("Stack trace:") self.echo.print(e.stacktrace) return False if not response.data.endswith('\n'): self.echo.error( "Error: Request failed: `SELECT version();` query failed.") return False version = response.data.strip().split('.') self.server_version = (int(version[0]), int(version[1]), int(version[2])) self.echo.success( "Connected to ClickHouse server v{0}.{1}.{2}.\n".format( *self.server_version)) return True def load_config(self): self.config = read_config() self.multiline = self.config.getboolean('main', 'multiline') self.format = self.format or self.config.get('main', 'format') self.format_stdin = self.format_stdin or self.config.get( 'main', 'format_stdin') self.show_formatted_query = self.config.getboolean( 'main', 'show_formatted_query') self.highlight = self.config.getboolean('main', 'highlight') self.highlight_output = self.config.getboolean('main', 'highlight_output') self.highlight_truecolor = self.config.getboolean( 'main', 'highlight_truecolor') and os.environ.get('COLORTERM') self.conn_timeout = self.config.getfloat('http', 'conn_timeout') self.conn_timeout_retry = self.config.getint('http', 'conn_timeout_retry') self.conn_timeout_retry_delay = self.config.getfloat( 'http', 'conn_timeout_retry_delay') self.host = self.host or self.config.get('defaults', 'host') or '127.0.0.1' self.port = self.port or self.config.get('defaults', 'port') or 8123 self.user = self.user or self.config.get('defaults', 'user') or 'default' self.password = self.password or self.config.get( 'defaults', 'password') or '' self.database = self.database or self.config.get('defaults', 'db') or 'default' config_settings = dict(self.config.items('settings')) arg_settings = self.settings config_settings.update(arg_settings) self.settings = config_settings self.echo.colors = self.highlight def run(self, query, data): self.load_config() if data or query is not None: self.format = self.format_stdin self.echo.verbose = False if self.echo.verbose: show_version() if not self.connect(): return if self.client: self.client.settings = self.settings self.client.cli_settings = { 'multiline': self.multiline, 'format': self.format, 'format_stdin': self.format_stdin, 'show_formatted_query': self.show_formatted_query, 'highlight': self.highlight, 'highlight_output': self.highlight_output, } if data and query is None: # cat stuff.sql | clickhouse-cli # clickhouse-cli stuff.sql for subdata in data: self.handle_input(subdata.read(), verbose=False, refresh_metadata=False) return if not data and query is not None: # clickhouse-cli -q 'SELECT 1' return self.handle_query(query, stream=True) if data and query is not None: # cat stuff.csv | clickhouse-cli -q 'INSERT INTO stuff' # clickhouse-cli -q 'INSERT INTO stuff' stuff.csv for subdata in data: compress = 'gzip' if os.path.splitext( subdata.name)[1] == '.gz' else False self.handle_query(query, data=subdata, stream=True, compress=compress) return layout = create_prompt_layout( lexer=PygmentsLexer(CHLexer) if self.highlight else None, get_prompt_tokens=get_prompt_tokens, get_continuation_tokens=get_continuation_tokens, multiline=self.multiline, ) buffer = CLIBuffer( client=self.client, multiline=self.multiline, metadata=self.metadata, ) application = Application( layout=layout, buffer=buffer, style=CHStyle if self.highlight else None, key_bindings_registry=KeyBinder.registry, ) eventloop = create_eventloop() self.cli = CommandLineInterface(application=application, eventloop=eventloop) self.cli.application.buffer.completer.refresh_metadata() try: while True: try: cli_input = self.cli.run(reset_current_buffer=True) self.handle_input(cli_input.text) except KeyboardInterrupt: # Attempt to terminate queries for query_id in self.query_ids: self.client.kill_query(query_id) self.echo.error("\nQuery was terminated.") finally: self.query_ids = [] except EOFError: self.echo.success("Bye.") def handle_input(self, input_data, verbose=True, refresh_metadata=True): force_pager = False if input_data.endswith( r'\p' if isinstance(input_data, str) else rb'\p'): input_data = input_data[:-2] force_pager = True # FIXME: A dirty dirty hack to make multiple queries (per one paste) work. self.query_ids = [] for query in sqlparse.split(input_data): query_id = str(uuid4()) self.query_ids.append(query_id) self.handle_query(query, verbose=verbose, query_id=query_id, force_pager=force_pager) if refresh_metadata and input_data: self.cli.application.buffer.completer.refresh_metadata() def handle_query(self, query, data=None, stream=False, verbose=False, query_id=None, compress=False, **kwargs): if query.rstrip(';') == '': return elif query.lower() in EXIT_COMMANDS: raise EOFError elif query.lower() in (r'\?', 'help'): rows = [ ['', ''], ["clickhouse-cli's custom commands:", ''], ['---------------------------------', ''], ['USE', "Change the current database."], ['SET', "Set an option for the current CLI session."], ['QUIT', "Exit clickhouse-cli."], ['HELP', "Show this help message."], ['', ''], ["PostgreSQL-like custom commands:", ''], ['--------------------------------', ''], [r'\l', "Show databases."], [r'\c', "Change the current database."], [r'\d, \dt', "Show tables in the current database."], [r'\d+', "Show table's schema."], [r'\ps', "Show current queries."], [r'\kill', "Kill query by its ID."], ['', ''], ["Query suffixes:", ''], ['---------------', ''], [r'\g, \G', "Use the Vertical format."], [r'\p', "Enable the pager."], ] for row in rows: self.echo.success('{:<8s}'.format(row[0]), nl=False) self.echo.info(row[1]) return elif query in (r'\d', r'\dt'): query = 'SHOW TABLES' elif query.startswith(r'\d+ '): query = 'DESCRIBE TABLE ' + query[4:] elif query == r'\l': query = 'SHOW DATABASES' elif query.startswith(r'\c '): query = 'USE ' + query[3:] elif query.startswith(r'\ps'): query = ( "SELECT query_id, user, address, elapsed, {}, memory_usage " "FROM system.processes WHERE query_id != '{}'").format( 'read_rows' if self.server_version[2] >= 54115 else 'rows_read', query_id) elif query.startswith(r'\kill '): self.client.kill_query(query[6:]) return response = '' self.progress_reset() try: response = self.client.query( query, fmt=self.format, data=data, stream=stream, verbose=verbose, query_id=query_id, compress=compress, ) except TimeoutError: self.echo.error("Error: Connection timeout.") return except ConnectionError: self.echo.error("Error: Failed to connect.") return except DBException as e: self.progress_reset() self.echo.error("\nQuery:") self.echo.error(query) self.echo.error("\n\nReceived exception from server:") self.echo.error(e.error) if self.stacktrace and e.stacktrace: self.echo.print("\nStack trace:") self.echo.print(e.stacktrace) self.echo.print('\nElapsed: {elapsed:.3f} sec.\n'.format( elapsed=e.response.elapsed.total_seconds())) return total_rows, total_bytes = self.progress_reset() self.echo.print() if stream: data = response.iter_lines() if hasattr( response, 'iter_lines') else response.data for line in data: print(line.decode('utf-8', 'ignore')) else: if response.data != '': print_func = print if self.config.getboolean('main', 'pager') or kwargs.pop( 'force_pager', False): print_func = self.echo.pager should_highlight_output = (verbose and self.highlight and self.highlight_output and response.format in PRETTY_FORMATS) formatter = TerminalFormatter() if self.highlight and self.highlight_truecolor: formatter = TerminalTrueColorFormatter( style=CHPygmentsStyle) if should_highlight_output: print_func( pygments.highlight(response.data, CHPrettyFormatLexer(), formatter)) else: print_func(response.data) if response.message != '': self.echo.print(response.message) self.echo.print() self.echo.success('Ok. ', nl=False) if response.rows is not None: self.echo.print('{rows_count} row{rows_plural} in set.'.format( rows_count=response.rows, rows_plural='s' if response.rows != 1 else '', ), end=' ') if self.config.getboolean( 'main', 'timing') and response.time_elapsed is not None: self.echo.print( 'Elapsed: {elapsed:.3f} sec. Processed: {rows} rows, {bytes} ({avg_rps} rows/s, {avg_bps}/s)' .format( elapsed=response.time_elapsed, rows=numberunit_fmt(total_rows), bytes=sizeof_fmt(total_bytes), avg_rps=numberunit_fmt(total_rows / max(response.time_elapsed, 0.001)), avg_bps=sizeof_fmt(total_bytes / max(response.time_elapsed, 0.001)), ), end='') self.echo.print('\n') def progress_update(self, line): if not self.config.getboolean('main', 'timing'): return # Parse X-ClickHouse-Progress header now = datetime.now() progress = json.loads(line[23:].decode().strip()) progress = { 'timestamp': now, 'read_rows': int(progress['read_rows']), 'total_rows': int(progress['total_rows']), 'read_bytes': int(progress['read_bytes']), } # Calculate percentage completed and format initial message progress['percents'] = int( (progress['read_rows'] / progress['total_rows']) * 100) if progress['total_rows'] > 0 else 0 message = 'Progress: {} rows, {}'.format( numberunit_fmt(progress['read_rows']), sizeof_fmt(progress['read_bytes'])) # Calculate row and byte read velocity if self.progress: delta = (now - self.progress['timestamp']).total_seconds() if delta > 0: rps = (progress['read_rows'] - self.progress['read_rows']) / delta bps = (progress['read_bytes'] - self.progress['read_bytes']) / delta message += ' ({} rows/s, {}/s)'.format(numberunit_fmt(rps), sizeof_fmt(bps)) self.progress = progress self.progress_print(message, progress['percents']) def progress_reset(self): progress = self.progress self.progress = None clickhouse_cli.helpers.trace_headers_stream = self.progress_update # Clear printed progress (if any) columns = shutil.get_terminal_size((80, 0)).columns sys.stdout.write(u"\u001b[%dD" % columns + " " * columns) sys.stdout.flush() # Report totals if progress: return (progress['read_rows'], progress['read_bytes']) return (0, 0) def progress_print(self, message, percents): suffix = '%3d%%' % percents columns = shutil.get_terminal_size((80, 0)).columns bars_max = columns - (len(message) + len(suffix) + 3) bars = int(percents * (bars_max / 100)) if (bars_max > 0) else 0 message = '{} \033[42m{}\033[0m{} {}'.format(message, " " * bars, " " * (bars_max - bars), suffix) sys.stdout.write(u"\u001b[%dD" % columns + message) sys.stdout.flush()
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) def connect(self): self.url = 'http://{host}:{port}/'.format(host=self.host, port=self.port) self.client = Client(self.url, self.user, self.password, self.database, self.settings, self.stacktrace) self.echo.print("Connecting to {host}:{port}".format(host=self.host, port=self.port)) try: response = self.client.query('SELECT version();', fmt='TabSeparated', timeout=10) 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_output = self.config.getboolean('main', 'highlight_output') self.host = self.host or self.config.get('defaults', 'host') or '127.0.0.1' self.port = self.port or self.config.get('defaults', 'port') or 8123 self.user = self.user or self.config.get('defaults', 'user') or 'default' self.database = self.database or self.config.get('defaults', 'db') or 'default' config_settings = dict(self.config.items('settings')) arg_settings = self.settings config_settings.update(arg_settings) self.settings = config_settings if self.client: self.client.settings = self.settings def run(self, query=None, data=None): self.load_config() if data is not None or query is not None: self.format = self.format_stdin self.echo.verbose = False if self.echo.verbose: show_version() if not self.connect(): return if data is not None and query is None: # cat stuff.sql | clickhouse-cli return self.handle_input('\n'.join(data), verbose=False) if data is None and query is not None: # clickhouse-cli -q 'SELECT 1' return self.handle_query(query, stream=True) if data is not None and query is not None: # cat stuff.csv | clickhouse-cli -q 'INSERT INTO stuff' return self.handle_query(query, data=data, stream=True) layout = create_prompt_layout( lexer=PygmentsLexer(CHLexer), get_prompt_tokens=get_prompt_tokens, get_continuation_tokens=get_continuation_tokens, multiline=self.multiline, ) buffer = CLIBuffer( client=self.client, multiline=self.multiline, ) application = Application( layout=layout, buffer=buffer, style=CHStyle, key_bindings_registry=KeyBinder.registry, ) eventloop = create_eventloop() cli = CommandLineInterface(application=application, eventloop=eventloop) try: while True: try: cli_input = cli.run(reset_current_buffer=True) self.handle_input(cli_input.text) except KeyboardInterrupt: # Attempt to terminate queries for query_id in self.query_ids: self.client.kill_query(query_id) self.echo.error("\nQuery was terminated.") finally: self.query_ids = [] except EOFError: self.echo.success("Bye.") def handle_input(self, input_data, verbose=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) def handle_query(self, query, data=None, stream=False, verbose=False, query_id=None): if query.rstrip(';') == '': return elif query.lower() in EXIT_COMMANDS: raise EOFError elif query.lower() in ('\?', '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:", ''], ['--------------------------------', ''], ['\l', "Show databases."], ['\c', "Change the current database."], ['\d, \dt', "Show tables in the current database."], ['\d+', "Show table's schema."], ['\ps', "Show current queries."], ['\kill', "Kill query by its ID."], ['', ''], ] for row in rows: self.echo.success('{:<8s}'.format(row[0]), nl=False) self.echo.info(row[1]) return elif query in ('\d', '\dt'): query = 'SHOW TABLES' elif query in ('\l',): query = 'SHOW DATABASES' elif query.startswith('\d+ '): query = 'DESCRIBE TABLE ' + query[4:] elif query.startswith('\c '): query = 'USE ' + query[3:] elif query.startswith('\ps'): if self.server_version[2] < 54115: query = "SELECT query_id, user, address, elapsed, rows_read, memory_usage FROM system.processes WHERE query_id != '{}'".format(query_id) else: query = "SELECT query_id, user, address, elapsed, read_rows, memory_usage FROM system.processes WHERE query_id != '{}'".format(query_id) elif query.startswith('\kill '): self.client.kill_query(query[6:]) return response = '' try: response = self.client.query( query, fmt=self.format, data=data, stream=stream, verbose=verbose, show_formatted=self.show_formatted_query, query_id=query_id ) except DBException as e: self.echo.error("\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 self.echo.print() if stream: print('\n'.join((e.decode('utf-8', 'ignore') for e in response.data)), end='') else: if response.data != '': print_func = self.echo.pager if self.config.getboolean('main', 'pager') else print if verbose and self.highlight_output and response.format in PRETTY_FORMATS: print_func(pygments.highlight( response.data, CHPrettyFormatLexer(), TerminalTrueColorFormatter(style=CHPygmentsStyle) )) 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.'.format( elapsed=response.time_elapsed ), end='') self.echo.print('\n')