stevekaplan123/carpe_diem

#!/usr/bin/env ruby

require 'pg' unless defined?( PG )

# 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?

        # This will be swapped soon for code that makes options like those required for
        # PQconnectdbParams()/PQconnectStartParams(). For now, stick to an options string for
        # PQconnectdb()/PQconnectStart().

        # 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 }
            appname = $0.sub(/^(.{30}).{4,}(.{30})$/){ $1+"..."+$2 }
            appname = PG::Connection.quote_connstr( appname )
            connopts = ["fallback_application_name=#{appname}"]
        else
            connopts = []
        end

        # Handle an options hash first
        if args.last.is_a?( Hash )
            opthash = args.pop
            opthash.each do |key, val|
                connopts.push( "%s=%s" % [key, PG::Connection.quote_connstr(val)] )
            end
        end

        # Option string style
        if args.length == 1 && args.first.to_s.index( '=' )
            connopts.unshift( args.first )

        # Append positional parameters
        else
            args.each_with_index do |val, i|
                next unless val # Skip nil placeholders

                key = CONNECT_ARGUMENT_ORDER[ i ] or
                    raise ArgumentError, "Extra positional parameter %d: %p" % [ i+1, val ]
                connopts.push( "%s=%s" % [key, PG::Connection.quote_connstr(val.to_s)] )
            end
        end

        return connopts.join(' ')
    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