Skip to content

Instantly share code, notes, and snippets.

@arika
Last active August 29, 2015 14:08
Show Gist options
  • Save arika/8f7640e8f1bb8b3c10f5 to your computer and use it in GitHub Desktop.
Save arika/8f7640e8f1bb8b3c10f5 to your computer and use it in GitHub Desktop.
run groonga command on pry
#!/usr/bin/env ruby
# encoding: utf-8
#
# Requirements:
# * groonga command
# * jq command <http://stedolan.github.io/jq/>
# * pry 0.10.x
#
# Configuration:
# * ~/.groonga-pryrc
ENV['PRYRC'] = '~/.groonga-pryrc'
require 'io/wait'
require 'json'
require 'pry'
class GroongaCompleter
table_flag = %w(
TABLE_NO_KEY TABLE_HASH_KEY TABLE_PAT_KEY TABLE_DAT_KEY KEY_WITH_SIS
)
column_flag = %w(
COLUMN_SCALAR COLUMN_VECTOR COLUMN_INDEX COMPRESS_ZLIB COMPRESS_LZO
WITH_SECTION WITH_WEIGHT WITH_POSITION
)
normalize_flag = %w(
NONE REMOVE_BLANK WITH_TYPES WITH_CHECKS REMOVE_TOKENIZED_DELIMITER
)
token_filter = %w(TokenFilterStopWord TokenFilterStem)
tokenize_flag = %w(NONE ENABLE_TOKENIZED_DELIMITER)
tokenize_mode = %w(ADD GET)
query_flag = %w(
ALLOW_PRAGMA ALLOW_COLUMN ALLOW_UPDATE ALLOW_LEADING_NOT NONE
)
log_level = %(EMERG ALERT CRIT error warning notice info debug)
COMMANDS = {
'cache_limit' => {
'--max' => nil,
},
'check' => {
'--obj' => :table,
},
'clearlock' => {
:arg => :table_column,
},
'column_create' => {
'--table' => :table,
'--name' => nil,
'--flags' => column_flag,
'--type' => :type_or_table,
'--source' => :source,
},
'column_list' => {
'--table' => :table,
},
'column_remove' => {
'--table' => :table,
'--name' => :column,
},
'column_rename' => {
'--table' => :table,
'--name' => :column,
'--new_name' => nil,
},
'define_selector' => {
'--name' => nil,
'--table' => :table,
'--match_columns' => :column,
'--query' => :table_column,
'--filter' => :expr,
'--scorer' => nil,
'--sortby' => :expr,
'--output_columns' => :expr,
'--offset' => nil,
'--limit' => nil,
'--drilldown' => :column,
'--drilldown_sortby' => :expr,
'--drilldown_output_columns' => :expr,
'--drilldown_offset' => nil,
'--drilldown_limit' => nil,
},
'defrag' => {
:arg => :table_column,
},
'delete' => {
'--table' => :table,
'--key' => nil,
'--id' => nil,
'--filter' => :expr,
},
'dump' => {
'--tables' => :table,
},
'load' => {
'--values' => nil,
'--table' => :table,
'--columns' => :column,
'--ifexists' => :expr,
'--input_type' => %w(JSON),
'--each' => :expr,
},
'log_level' => {
'--level' => log_level,
},
'log_put' => {
'--level' => log_level,
'--message' => nil,
},
'log_reopen' => {},
'normalize' => {
'--normalizer' => :normalizer,
'--string' => nil,
'--flags' => normalize_flag,
},
'normalizer_list' => {},
'quit' => {},
'range_filter' => {}, # FIXME
'register' => {
'--path' => :plugin_path,
},
'ruby_eval' => {
'--script' => nil,
},
'ruby_load' => { # OK?
'--path' => :path
},
'select' => {
'--table' => :table,
'--match_columns' => :column,
'--query' => :table_column,
'--filter' => :expr,
'--scorer' => :expr,
'--sortby' => :expr,
'--output_columns' => :expr,
'--offset' => nil,
'--limit' => nil,
'--drilldown' => :column,
'--drilldown_sortby' => :expr,
'--drilldown_output_columns' => :expr,
'--drilldown_offset' => nil,
'--drilldown_limit' => nil,
'--cache' => %w(no),
'--match_escalation_threshold' => nil,
#'--query_expansion' => nil, # deprecated
'--query_flags' => query_flag,
'--query_expander' => :table_column,
'--adjuster' => :column,
},
'shutdown' => {},
'status' => {},
'suggest' => {
'--types' => %w(complete correct suggest),
'--table' => :item_table,
'--column' => :column, # XXX: OK?
'--query' => :table_column,
'--sortby' => :expr,
'--output_columns' => :expr,
'--offset' => nil,
'--limit' => nil,
'--frequency_threshold' => nil,
'--conditional_probability_threshold' => nil,
'--prefix_search' => %w(yes no auto),
'--similar_search' => %w(yes no auto),
},
'table_create' => {
'--name' => nil,
'--flags' => table_flag,
'--key_type' => :type_or_table,
'--value_type' => :type, # XXX: OK?
'--default_tokenizer' => :torkenizer,
'--normalizer' => :normalizer,
'--token_filters' => token_filter,
},
'table_list' => {},
'table_remove' => {
'--name' => :table,
},
'table_tokenize' => {
'--table' => :table,
'--string' => nil,
'--flags' => tokenize_flag,
'--mode' => tokenize_mode,
},
'tokenize' => {
'--tokenizer' => :tokenizer,
'--string' => nil,
'--normalizer' => :normalizer,
'--flags' => tokenize_flag,
'--mode' => tokenize_mode,
'--token_filters' => token_filter,
},
'tokenizer_list' => {},
'truncate' => {
'--table_name' => :table,
},
}
OPTIONS = COMMANDS.values.map(&:keys).flatten.uniq.grep(/\A--/)
DATA_TYPES = %w(
Object Bool Int8 UInt8 Int16 UInt16 Int32 UInt32 Int64 UInt64 Float Time
ShortText Text LongText TokyoGeoPoint WGS84GeoPoint
)
PSEUDO_COLUMNS = %w(
_id _key _value _score _nsubrecs
)
FUNCTIONS = %w(
between edit_distance geo_distance geo_in_circle geo_in_rectangle
highlight_full highlight_html html_untag in_values now query rand
snippet_html sub_filter
).map {|fn| "#{fn}()" }
COMMAND_REGEXP = /(#{COMMANDS.keys.map {|cmd| Regexp.quote(cmd) }.join('|')})/
@@orig_completer = nil
def self.orig_completer=(completer)
@@orig_completer = completer
end
def initialize(input, pry = nil)
@input = input
@pry = pry
@orig_instance = @@orig_completer.new(@input, @pry)
end
def call(str, options = {})
return commands if empty_input?
return @orig_instance.call(str, options) if not_a_command?
prev_idx, idx, args = split_command_line
*str_prefixs, str_key = str.split(/[,]/)
if str_prefixs.empty?
str_prefix = ''
else
str_prefix = str_prefixs.join(',') + ','
end
if idx == 0
return filter(commands, str_key, str_prefix) +
@orig_instance.call(str, options)
end
if prev_idx.nil? || /\A--/ !~ args[prev_idx]
return filter(command_options(args.first) - args.grep(/\A--/), str_key, str_prefix)
end
list = option_candidate(prev_idx, idx, args, str_key)
filter(list, str_key, str_prefix)
end
private
def argument_value(args, name)
if i = args.index(name)
args[i + 1]
else
nil
end
end
def option_candidate(prev_idx, idx, args, str_key)
arg_type = option_argument(args.first, args[prev_idx])
case arg_type
when nil
list = []
when Array
list = arg_type
when :table
list = get_table_list
when :table_column
list = get_table_column_list(str_key)
when :column
table = argument_value(args, '--table')
list = get_column_list(table)
when :type_or_table
list = get_table_list + DATA_TYPES
when :source
table = argument_value(args, '--type')
list = get_column_list(table)
when :expr
list = (get_table_column_list(str_key) +
PSEUDO_COLUMNS + FUNCTIONS).uniq
when :normalizer
list = get_list('normalizer_list', 'name')
when :tokenizer
list = get_list('tokenizer_list', 'name')
when :plugin_path
file_completer_proc = @input::FILENAME_COMPLETION_PROC
if /\A\// =~ str_key
list = file_completer_proc.call(str_key || '') || []
else
plugins_dir = `pkg-config --variable=pluginsdir groonga`.chomp
path_key = File.join(plugins_dir, str_key || '')
list = (file_completer_proc.call(path_key) || []).
select do |path|
/\.so\z/ =~ path || FileTest.directory?(path)
end.map do |path|
path[plugins_dir.size + 1 .. -1].sub(/\.so\z/, '')
end
end
when :path
list = @input::FILENAME_COMPLETION_PROC.call(str_key || '') || []
when :item_table
list = %w(item_)
else
list = OPTIONS
end
list
end
def get_list(command, target)
groonga = @pry.config.groonga_process
groonga.write "#{command}\n"
stat, rows = groonga.read_result
if stat[0] == 0
rows.map {|row| row[target] }
else
[]
end
end
def get_schema(command, target)
groonga = @pry.config.groonga_process
groonga.write "#{command}\n"
stat, (schema, *rows) = groonga.read_result
if stat[0] == 0 &&
i = schema.index {|col_name, col_type| col_name == target }
rows.map {|row| row[i] }
else
[]
end
end
def get_table_list
get_schema('table_list', 'name')
end
def get_column_list(table)
if table
list = get_schema("column_list #{table}", 'name')
else
list = get_table_list.map do |table|
get_column_list(table)
end.flatten
end
(list + PSEUDO_COLUMNS).uniq
end
def get_table_column_list(str_key)
if /\./ =~ (str_key || '')
table, column = $`, $'
get_column_list(table).map {|c| "#{table}.#{c}" }
else
get_table_list
end
end
def commands
COMMANDS.keys
end
def command_options(command)
return [] unless COMMANDS.include?(command)
COMMANDS[command].keys
end
def option_argument(command, name)
return {} unless COMMANDS.include?(command)
args = COMMANDS[command]
return {} unless args.include?(name)
args[name]
end
def line_buffer
@input.line_buffer
end
def empty_input?
line_buffer.empty?
end
def not_a_command?
shell_words = line_buffer.strip.split(/\s+/)
shell_words.size > 1 &&
/\A#{COMMAND_REGEXP}\z/o !~ shell_words.first
end
def split_command_line
point = @input.point
args = []
prev_idx = idx = nil
buf = nil
len = 0
segs = shell_split(line_buffer)
segs.each do |seg|
if /\A\s+\z/ =~ seg
args << buf if buf
buf = nil
else
buf ||= ''
buf << seg
end
len += seg.size
if prev_idx.nil? && idx.nil? && len >= point
idx = args.size if buf # current shell-word is pointed
prev_idx = args.size - 1 unless args.empty?
end
end
args << buf if buf
if prev_idx.nil? && idx.nil?
prev_idx = args.size - 1
end
[prev_idx, idx, args]
end
def shell_split(line)
segs = []
pos = 0
patt = /((["'])(?:\\.|(?!\2)[^\\])*\2)|(?:\\.|[^\\'"\s])+|\s+/o
while line.match(patt, pos)
break unless $~.begin(0) == pos
pos = $~.end(0)
segs << $&
end
unless pos == line.size
segs << line[pos .. -1]
end
segs
end
def filter(list, key, prefix)
return list unless key
key_regexp = /\A#{Regexp.quote(key)}/i
filtered = list.grep(key_regexp)
if prefix
filtered.map {|k| prefix + k }
else
filtered
end
end
end
class GroongaProcess
def initialize(argv)
@argv = argv
@pid = nil
@last_result = nil
start
end
attr_reader :last_result, :pid
def write(data)
@to_grn.print data
rescue Errno::EPIPE, IOError => e
close
raise e
end
def read
output = ''
begin
loop do
output << @fr_grn.readpartial(1024)
break unless ready?
end
rescue EOFError
close
end
output
rescue Errno::EPIPE, IOError => e
close
raise e
end
def read_result
output = read
result = JSON.parse(output)
result.define_singleton_method(:to_s) do
output
end
@last_result = result
rescue JSON::ParserError
raise "unexpected result: #{output}"
end
def ready?
@fr_grn.ready?
end
def close
@to_grn.close unless @to_grn.closed?
@fr_grn.close unless @fr_grn.closed?
@pid = nil
ensure
Process.wait(@pid) rescue Errno::ECHILD
end
def kill
return unless @pid
Process.kill('TERM', @pid)
Process.wait(@pid)
begin
Process.kill(0, @pid)
raise "Could not stop current process: #{@pid}"
rescue Errno::ESRCH
@pid = nil
end
end
def restart
kill
start
end
private
def start
to_grn0, @to_grn = IO.pipe
@fr_grn, fr_grn0 = IO.pipe
@pid = Process.spawn(
'groonga',
'--input-fd', to_grn0.fileno.to_s,
'--output-fd', fr_grn0.fileno.to_s,
*@argv,
to_grn0 => to_grn0, fr_grn0 => fr_grn0,
pgroup: true)
to_grn0.close
fr_grn0.close
end
end
Pry::Commands.create_command GroongaCompleter::COMMAND_REGEXP do
description 'Run command.'
banner <<-BANNER
Usage: <groonga-command> [groonga-options...]
BANNER
command_options(
listing: 'groonga-command',
keep_retval: true,
takes_block: true,
shellwords: false,
requires_gem: 'json',
)
def process
retval = run_command
if command_block
retval = command_block.call(retval)
end
retval
rescue Errno::EPIPE, IOError
output.puts "#{$!.class}: #{$!.message}"
void
end
private
def run_command
data = slice_heredoc_data!
if /\A<<-?(\w+)\z/ =~ args.last
args.pop
data = read_heredoc_data($1)
end
if data.nil? && args.first == 'load'
data = read_stdin_data
end
groonga = Pry.config.groonga_process
groonga.write args.join(' ') + "\n"
if data
groonga.write data
end
if args.first == 'dump'
groonga.read
else
groonga.read_result
end
end
def read_stdin_data
data = ''
while line = $stdin.gets
data << line
end
data
end
def read_heredoc_data(eoh)
eoh_line = "#{eoh}\n"
data = ''
while line = $stdin.gets
break if line == eoh_line
data << line
end
data
end
def slice_heredoc_data!
return if eval_string.empty?
return unless /\A[ \t]*(?:['"]|%[qQ]?[\x20-\x2f\x3a-\x40]|<<-?\w+[ \t]*\n)/ =~ eval_string
eval_string.replace('')
$'.strip + "\n"
end
end
Pry::Commands.create_command 'kill-groonga' do
description 'Kill groonga process'
banner <<-BANNER
Usage: kill-groonga
BANNER
def process
output.puts Pry.config.groonga_process.kill
rescue Errno::ESRCH
output.puts "#{$!.class}: #{$!.message}"
end
end
Pry::Commands.create_command 'restart-groonga' do
description 'Restart groonga process'
banner <<-BANNER
Usage: restart-groonga
BANNER
def process
output.puts Pry.config.groonga_process.restart
end
end
Pry::Commands.create_command 'jq' do
description 'Invoke jq for groonga results.'
banner <<-BANNER
Usage: jq [[json-text] filter]
BANNER
command_options(
shellwords: true
)
def process
if args.size < 2
groonga = Pry.config.groonga_process
text = groonga.last_result
filter = args.first
else
text, filter, = args
end
_jq_print(text, filter, _pry_) if text
end
end
def _jq_print(json, filter, pry_instance)
to_jq0, to_jq = IO.pipe
fr_jq, fr_jq0 = IO.pipe
er_jq, er_jq0 = IO.pipe
color_opt = Pry.config.color ? '-C' : '-M'
pid = Process.spawn('jq', color_opt, filter || '.',
in: to_jq0, out: fr_jq0, err: er_jq0)
to_jq0.close
fr_jq0.close
er_jq0.close
to_jq.write json
to_jq.close
err = er_jq.read
er_jq.close
if err.empty?
pry_instance.pager.page fr_jq.read
fr_jq.close
else
pry_instance.output.print "jq: #{err}"
end
ensure
Process.wait(pid) if pid
end
grn_process = GroongaProcess.new(ARGV)
begin
Pry.config.should_load_local_rc = false
Pry.config.groonga_process = grn_process
Pry.config.history.file = '~/.groonga-pry_history'
Pry.config.hooks.add_hook(:when_started, :setup_completer) do |target, opts, pry_instance|
GroongaCompleter.orig_completer = pry_instance.config.completer
pry_instance.config.completer = GroongaCompleter
end
Pry.config.hooks.add_hook(:when_started, :read_initial_result) do |target, opts, pry_instance|
define_method(:jq) do |json, filter = nil|
_jq_print(json, filter, pry_instance)
end
if grn_process.ready?
grn_out = grn_process.read
pry_instance.output.print grn_out
end
end
binding.pry(quiet: true, prompt_name: 'groonga')
ensure
grn_process.close if grn_process
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment