#!/usr/bin/env ruby

require 'pg' unless defined?( PG )
require 'uri'

# The PostgreSQL connection class. The interface for this class is based on
# {libpq}[http://www.postgresql.org/docs/9.2/interactive/libpq.html], the C
# application programmer's interface to PostgreSQL. Some familiarity with libpq
# is recommended, but not necessary.
#
# For example, to send query to the database on the localhost:
#
#    require 'pg'
#    conn = PG::Connection.open(:dbname => 'test')
#    res = conn.exec_params('SELECT $1 AS a, $2 AS b, $3 AS c', [1, 2, nil])
#    # Equivalent to:
#    #  res  = conn.exec('SELECT 1 AS a, 2 AS b, NULL AS c')
#
# See the PG::Result class for information on working with the results of a query.
#
class PG::Connection

	# The order the options are passed to the ::connect method.
	CONNECT_ARGUMENT_ORDER = %w[host port options tty dbname user password]


	### Quote the given +value+ for use in a connection-parameter string.
	def self::quote_connstr( value )
		return "'" + value.to_s.gsub( /[\\']/ ) {|m| '\\' + m } + "'"
	end


	### Parse the connection +args+ into a connection-parameter string. See PG::Connection.new
	### for valid arguments.
	def self::parse_connect_args( *args )
		return '' if args.empty?

		hash_arg = args.last.is_a?( Hash ) ? args.pop : {}
		option_string = ''
		options = {}

		# Parameter 'fallback_application_name' was introduced in PostgreSQL 9.0
		# together with PQescapeLiteral().
		if PG::Connection.instance_methods.find {|m| m.to_sym == :escape_literal }
			options[:fallback_application_name] = $0.sub( /^(.{30}).{4,}(.{30})$/ ){ $1+"..."+$2 }
		end

		if args.length == 1
			case args.first
			when URI, URI.regexp
				uri = URI(args.first)
				options.merge!( Hash[URI.decode_www_form( uri.query )] ) if uri.query
			when /=/
				# Option string style
				option_string = args.first.to_s
			else
				# Positional parameters
				options[CONNECT_ARGUMENT_ORDER.first.to_sym] = args.first
			end
		else
			max = CONNECT_ARGUMENT_ORDER.length
			raise ArgumentError,
				"Extra positional parameter %d: %p" % [ max + 1, args[max] ] if args.length > max

			CONNECT_ARGUMENT_ORDER.zip( args ) do |(k,v)|
				options[ k.to_sym ] = v if v
			end
		end

		options.merge!( hash_arg )

		if uri
			uri.host     = nil if options[:host]
			uri.port     = nil if options[:port]
			uri.user     = nil if options[:user]
			uri.password = nil if options[:password]
			uri.path     = '' if options[:dbname]
			uri.query    = URI.encode_www_form( options )
			return uri.to_s.sub( /^#{uri.scheme}:(?!\/\/)/, "#{uri.scheme}://" )
		else
			option_string += ' ' unless option_string.empty? && options.empty?
			return option_string + options.map { |k,v| "#{k}=#{quote_connstr(v)}" }.join( ' ' )
		end
	end


	#  call-seq:
	#     conn.copy_data( sql ) {|sql_result| ... } -> PG::Result
	#
	# Execute a copy process for transfering data to or from the server.
	#
	# This issues the SQL COPY command via #exec. The response to this
	# (if there is no error in the command) is a PG::Result object that
	# is passed to the block, bearing a status code of PGRES_COPY_OUT or
	# PGRES_COPY_IN (depending on the specified copy direction).
	# The application should then use #put_copy_data or #get_copy_data
	# to receive or transmit data rows and should return from the block
	# when finished.
	#
	# #copy_data returns another PG::Result object when the data transfer
	# is complete. An exception is raised if some problem was encountered,
	# so it isn't required to make use of any of them.
	# At this point further SQL commands can be issued via #exec.
	# (It is not possible to execute other SQL commands using the same
	# connection while the COPY operation is in progress.)
	#
	# This method ensures, that the copy process is properly terminated
	# in case of client side or server side failures. Therefore, in case
	# of blocking mode of operation, #copy_data is preferred to raw calls
	# of #put_copy_data, #get_copy_data and #put_copy_end.
	#
	# Example with CSV input format:
	#   conn.exec "create table my_table (a text,b text,c text,d text,e text)"
	#   conn.copy_data "COPY my_table FROM STDIN CSV" do
	#     conn.put_copy_data "some,csv,data,to,copy\n"
	#     conn.put_copy_data "more,csv,data,to,copy\n"
	#   end
	# This creates +my_table+ and inserts two rows.
	#
	# Example with CSV output format:
	#   conn.copy_data "COPY my_table TO STDOUT CSV" do
	#     while row=conn.get_copy_data
	#       p row
	#     end
	#   end
	# This prints all rows of +my_table+ to stdout:
	#   "some,csv,data,to,copy\n"
	#   "more,csv,data,to,copy\n"
	def copy_data( sql, coder=nil )
		res = exec( sql )

		case res.result_status
		when PGRES_COPY_IN
			begin
				if coder
					old_coder = self.encoder_for_put_copy_data
					self.encoder_for_put_copy_data = coder
				end
				yield res
			rescue Exception => err
				errmsg = "%s while copy data: %s" % [ err.class.name, err.message ]
				put_copy_end( errmsg )
				get_result
				raise
			else
				put_copy_end
				get_last_result
			ensure
				self.encoder_for_put_copy_data = old_coder if coder
			end

		when PGRES_COPY_OUT
			begin
				if coder
					old_coder = self.decoder_for_get_copy_data
					self.decoder_for_get_copy_data = coder
				end
				yield res
			rescue Exception => err
				cancel
				while get_copy_data
				end
				while get_result
				end
				raise
			else
				res = get_last_result
				if res.result_status != PGRES_COMMAND_OK
					while get_copy_data
					end
					while get_result
					end
					raise PG::NotAllCopyDataRetrieved, "Not all COPY data retrieved"
				end
				res
			ensure
				self.decoder_for_get_copy_data = old_coder if coder
			end

		else
			raise ArgumentError, "SQL command is no COPY statement: #{sql}"
		end
	end

	# Backward-compatibility aliases for stuff that's moved into PG.
	class << self
		define_method( :isthreadsafe, &PG.method(:isthreadsafe) )
	end


	### Returns an array of Hashes with connection defaults. See ::conndefaults
	### for details.
	def conndefaults
		return self.class.conndefaults
	end

	### Return the Postgres connection defaults structure as a Hash keyed by option
	### keyword (as a Symbol).
	###
	### See also #conndefaults
	def self.conndefaults_hash
		return self.conndefaults.each_with_object({}) do |info, hash|
			hash[ info[:keyword].to_sym ] = info[:val]
		end
	end

	### Returns a Hash with connection defaults. See ::conndefaults_hash
	### for details.
	def conndefaults_hash
		return self.class.conndefaults_hash
	end

	# Method 'conninfo' was introduced in PostgreSQL 9.3.
	if self.instance_methods.find{|m| m.to_sym == :conninfo }

		### Return the Postgres connection info structure as a Hash keyed by option
		### keyword (as a Symbol).
		###
		### See also #conninfo
		def conninfo_hash
			return self.conninfo.each_with_object({}) do |info, hash|
				hash[ info[:keyword].to_sym ] = info[:val]
			end
		end
	end

end # class PG::Connection

# Backward-compatible alias
PGconn = PG::Connection