module PgCsv::Base

Public Class Methods

new(opts = {}) click to toggle source
# File lib/pg_csv.rb, line 8
def initialize(opts = {})
  @options = opts
end

Public Instance Methods

export(to = nil, opts = {}, &row_proc) click to toggle source

do export :to - filename or stream

# File lib/pg_csv.rb, line 13
def export(to = nil, opts = {}, &row_proc)
  @row_proc = row_proc
  @local_options = opts

  raise ":connection should be" unless connection
  raise ":sql should be" unless sql

  with_temp_file?(to, temp_file, temp_dir) do |dest|
    export_to(dest)
  end
end

Protected Instance Methods

check_str(to) click to toggle source
# File lib/pg_csv.rb, line 81
def check_str(to)
  raise "'to' should be an string" unless to.is_a?(String)
end
columns_str() click to toggle source
# File lib/pg_csv.rb, line 166
def columns_str
  if o(:columns)
    col = o(:columns)
    if col.is_a?(Array)
      col.join(delimiter) + "\n"
    else
      col + "\n"
    end
  end
end
connection() click to toggle source
# File lib/pg_csv.rb, line 150
def connection
  o(:connection) || (defined?(ActiveRecord::Base) ? ActiveRecord::Base.connection : nil)
end
delimiter() click to toggle source
# File lib/pg_csv.rb, line 177
def delimiter
  o(:delimiter) || ','
end
encoding() click to toggle source
# File lib/pg_csv.rb, line 193
def encoding
  o(:encoding)
end
export_to(to) click to toggle source
# File lib/pg_csv.rb, line 41
def export_to(to)

  start = Time.now
  info "===> start generate export #{to}, type: #{type}"

  result = nil
  exporter = method(:export_to_stream).to_proc

  case type

    when :file
      check_str(to)
      File.open(to, 'w', &exporter)

    when :gzip
      check_str(to)
      require 'zlib'
      ::Zlib::GzipWriter.open(to, &exporter)

    when :stream
      raise "'to' should be" unless to
      exporter[to]

    when :plain
      require 'stringio'
      sio = StringIO.new
      exporter[sio]
      result = sio.string

    when :yield
      # not real saving anywhere, just yield each record
      raise "row_proc should be" unless @row_proc
      result = load_data{|_|}
  end

  info "<=== finished write #{to} in #{Time.now - start}"

  result
end
export_to_stream(stream) click to toggle source
# File lib/pg_csv.rb, line 85
def export_to_stream(stream)
  count = write_csv(stream)
  stream.flush if stream.respond_to?(:flush) && count > 0

  info "<= done exporting (#{count}) records."
end
force_quote() click to toggle source
# File lib/pg_csv.rb, line 197
def force_quote
  o(:force_quote)
end
info(message) click to toggle source
# File lib/pg_csv.rb, line 140
def info(message)
  logger.info(message) if logger
end
load_data() { |row_proc ? row_proc : columns_str| ... } click to toggle source
# File lib/pg_csv.rb, line 98
def load_data
  info "#{query}"
  raw = connection.respond_to?(:raw_connection) ? connection.raw_connection : connection
  count = 0

  info "=> query"
  raw.copy_data(query) do
    info "<= query"

    info "=> write data"
    if columns_str
      yield(@row_proc ? @row_proc[columns_str] : columns_str)
    end

    if @row_proc
      while row = raw.get_copy_data()
        yield(@row_proc[row])
        count += 1
      end
    else
      while row = raw.get_copy_data()
        yield(row)
        count += 1
      end
    end
    info "<= write data"

  end
  count
end
logger() click to toggle source
# File lib/pg_csv.rb, line 154
def logger
  o(:logger)
end
o(key) click to toggle source

options/defaults =============

# File lib/pg_csv.rb, line 146
def o(key)
  @local_options[key] || @options[key]
end
query() click to toggle source
# File lib/pg_csv.rb, line 129
    def query
      <<-SQL
  COPY (
    #{sql}
  ) TO STDOUT
  WITH CSV
  DELIMITER '#{delimiter}'
  #{use_pg_header? ? 'HEADER' : ''} #{encoding ? "ENCODING '#{encoding}'" : ''} #{force_quote ? "FORCE QUOTE *" : ''}
      SQL
    end
sql() click to toggle source
# File lib/pg_csv.rb, line 181
def sql
  o(:sql)
end
temp_dir() click to toggle source
# File lib/pg_csv.rb, line 189
def temp_dir
  o(:temp_dir) || '/tmp'
end
temp_file() click to toggle source
# File lib/pg_csv.rb, line 185
def temp_file
  o(:temp_file)
end
type() click to toggle source
# File lib/pg_csv.rb, line 158
def type
  o(:type) || :file
end
use_pg_header?() click to toggle source
# File lib/pg_csv.rb, line 162
def use_pg_header?
  o(:header) && !o(:columns)
end
with_temp_file?(to, use_temp_file, tmp_dir) { |filename| ... } click to toggle source
# File lib/pg_csv.rb, line 27
def with_temp_file?(to, use_temp_file, tmp_dir)
  if use_temp_file && [:file, :gzip].include?(type)
    check_str(to)

    self.class.with_temp_file(to, tmp_dir) do |filename|
      yield(filename)
    end

    info "<=== moving export to #{to}"
  else
    yield(to)
  end
end
write_csv(stream) click to toggle source
# File lib/pg_csv.rb, line 92
def write_csv(stream)
  load_data do |row|
    stream.write(row)
  end
end