#!/usr/bin/env ruby

########################################################################
# rq - Simple command-line MySQL query client                          #
# by Paul Duncan <pabs@pablotron.org>                                  #
#                                                                      #
# Copyright (C) 2002 Paul Duncan                                       #
#                                                                      #
# Permission is hereby granted, free of charge, to any person          #
# obtaining a copy of this software and associated documentation files #
# (the "Software"), to deal in the Software without restriction,       #
# including without limitation the rights to use, copy, modify, merge, #
# publish, distribute, sublicense, and/or sell copies of the Software, #
# and to permit persons to whom the Software is furnished to do so,    #
# subject to the following conditions:                                 #
#                                                                      #
# The above copyright notice and this permission notice shall be       #
# included in all copies of the Software, its documentation and        #
# marketing & publicity materials, and acknowledgment shall be given   #
# in the documentation, materials and software packages that this      #
# Software was used.                                                   #
#                                                                      #
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,      #
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF   #
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND                #
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY     #
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, #
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE    #
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.               #
########################################################################

require 'getoptlong'
require 'mysql'

# version number
RQ_VERSION = '0.1.3'
# path to default config file (only loaded if it exists)
DEFAULT_CONFIG = ENV['HOME'] + '/.rqrc'
# default output delimiter
DEFAULT_DELIMITER = ','

# usage function (called via --help command-line option)
def print_usage
  puts <<USAGE
rq - Command-line MySQL query tool.
Version #{RQ_VERSION} by Paul Duncan <pabs@pablotron.org>

Usage:
  rq [options] -s server -u user -p password -q query

Options:
  -c, --config   : Load configuration file (optional).
  -d, --database : Database to connect to (optional).
  -e, --delim    : Output field delimiter (optional, defaults to ',').
  -f, --force    : Allow world-writable config file (optional, DANGEROUS!).
  -h, --help     : Print out this help screen and exit.
  -m, --munge    : Munge delimiter to specified string in fields (optional).
  -n, --noconfig : Don't load default config file (optional).
  -q, --query    : Query string (required).
  -p, --password : Password (required).
  -s, --server   : MySQL server (required).
  -t, --notitles : Don't print field titles (optional).
  -u, --user     : Username (required).
  -v, --version  : Print out version information and exit.

Examples:
  # standard use: get all the records from a database
  #$0 -s localhost -u user -p pass -q 'select * from db.table'

  # read alternate config file
  #$0 -c sql_login.cfg --query 'select * from db.table'

  # convert commas in fields to underscores (db, user, and pass stored in
  # default config file)
  #$0 -q 'select * from db.table' -m _

Notes:
  * All options may be specified in the config file.  See the included
    _rqrc for sample use.
  * Field delimiters and munge strings longer than one character will
    produce unexpected results.
  * Newlines and binary data are not escaped when printing field contents.
  * rq will print an error and exit if your ~/.rqrc if it is world-
    writable (use the --force option to override this behavior).  Note
    that the rq config file is a standard ruby script so using this 
    option is a HUGE SECURITY HOLE and is NOT RECOMMENDED.

About the Author:
Paul Duncan <pabs@pablotron.org>
http://www.pablotron.org/
USAGE
  exit 0
end

def print_version
  puts "rq version #{RQ_VERSION}."
  exit 0
end

# perl does this right :)
def die(str)
  if str =~ /:$/
    $stderr.puts "ERROR: #{str} #{$!}."
  else
    $stderr.puts "ERROR: #{str}"
  end

  exit -1
end

# add csv_escape and csv_munge methods
class String
  def csv_escape(delim)
    gsub(/([#{delim}\\])/, "\\\\\\1")
  end

  def csv_munge(delim, munge)
    gsub(/#{delim}/, munge)
  end
end

# add to_hash method to GetoptLong
# (I proposed this as a standard feature on ruby-talk
# and received a less than spectacular response :-( )
class GetoptLong
  def to_hash
    hash = {}
    each { |key, val| hash[key] = val or true }
    hash
  end
end

#######################################################################

if ARGV.size > 1 or ARGV[0] =~ /^-/
  # parse command-line arguments and dump specified ones into hash
  $opts = GetoptLong.new(
    [ '--config',   '-c', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--database', '-d', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--delim',    '-e', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--force',    '-f', GetoptLong::NO_ARGUMENT ],
    [ '--help',     '-h', GetoptLong::NO_ARGUMENT ],
    [ '--munge',    '-m', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--noconfig', '-n', GetoptLong::NO_ARGUMENT ],
    [ '--notitles', '-t', GetoptLong::NO_ARGUMENT ],
    [ '--query',    '-q', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--password', '-p', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--server',   '-s', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--user',     '-u', GetoptLong::REQUIRED_ARGUMENT ],
    [ '--version',  '-v', GetoptLong::NO_ARGUMENT ]
  ).to_hash
else
  $opts = {}

  if ARGV.size == 1 and ARGV[0] !~ /^-/
    # if there's only one command-line argument, then make it the query
    # (assume we're getting the rest of the values from our ~/.rqrc)
    $opts['--query'] = ARGV[0].dup
  end
end

# handle specified command-line arguments
print_usage if $opts['--help']
print_version if $opts['--version']

# save a copy before loading the config file (so we can override the config
# file options with those specified on the command-line)
cmd_opts = $opts.dup;

# load a config file (if one was specified on the command-line, then use
# that. otherwise, check to see if the default config exists, and if so,
# use that instead)
if $opts['--config']
  # load specified config file
  Kernel::load $opts['--config'], false

  # replace whatever config file sets with command-line options
  cmd_opts.each { |key, val| $opts[key] = val }
elsif Kernel::test ?e, DEFAULT_CONFIG and !$opts['--noconfig']
  if !$opts['--force'] && File::stat(DEFAULT_CONFIG).mode & 02 != 0
    die "You have a world-writable \"#{DEFAULT_CONFIG}\".",
        'Please fix or remove this file before using rq.'
  end

  # load default config file
  Kernel::load DEFAULT_CONFIG, false
  
  # replace whatever config file sets with command-line options
  cmd_opts.each { |key, val| $opts[key] = val }
end

# check for required options
%w{--server --user --password --query}.each { |option|
  unless $opts[option]
    die "Missing required option \"#{option}\": see --help for usage."
  end
}
  
# set up output delimiter
$delim = $opts['--delim'] || DEFAULT_DELIMITER
die 'Only one-character delimiter strings are allowed.' if $delim.length > 1

# build list of Mysql::connect arguments
connect_args = [$opts['--server'], $opts['--user'], $opts['--password']]
connect_args << $opts['--database'] if $opts['--database']

# connect to db and make query
begin
  db = Mysql::connect *connect_args
  result = db.query $opts['--query']
rescue 
  die 'MySQL Error:'
end

# insert and delete don't return results
if result
  # print out field titles
  unless $opts['--notitles']
    puts (result.fetch_fields.map { |field|
      if $opts['--munge']
        field.name ? field.name.csv_munge($delim, $opts['--munge']) : ''
      else
        field.name ? field.name.csv_escape($delim) : ''
      end
    }.join $delim)
  end
  
  # print out each record
  result.each { |record|
    puts record.map { |field|
      if $opts['--munge']
        field ? field.csv_munge($delim, $opts['--munge']) : ''
      else
        field ? field.csv_escape($delim) : ''
      end
    }.join $delim
  }
end

# close the connection to the db
db.close