run groonga command on pry

groonga-pry.rb

#!/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