module PgSlice::Helpers

Constants

SQL_FORMAT

Protected Instance Methods

abort(message) click to toggle source
# File lib/pgslice/helpers.rb, line 21
def abort(message)
  raise Thor::Error, message
end
advance_date(date, period, count = 1) click to toggle source
# File lib/pgslice/helpers.rb, line 139
def advance_date(date, period, count = 1)
  date = date.to_date
  case period.to_sym
  when :day
    date.next_day(count)
  when :month
    date.next_month(count)
  else
    date.next_year(count)
  end
end
assert_no_table(table) click to toggle source
# File lib/pgslice/helpers.rb, line 135
def assert_no_table(table)
  abort "Table already exists: #{table}" if table.exists?
end
assert_table(table) click to toggle source
# File lib/pgslice/helpers.rb, line 131
def assert_table(table)
  abort "Table not found: #{table}" unless table.exists?
end
connection() click to toggle source

database connection

# File lib/pgslice/helpers.rb, line 27
def connection
  @connection ||= begin
    url = options[:url] || ENV["PGSLICE_URL"]
    abort "Set PGSLICE_URL or use the --url option" unless url

    uri = URI.parse(url)
    params = CGI.parse(uri.query.to_s)
    # remove schema
    @schema = Array(params.delete("schema") || "public")[0]
    uri.query = URI.encode_www_form(params)

    ENV["PGCONNECT_TIMEOUT"] ||= "1"
    conn = PG::Connection.new(uri.to_s)
    conn.set_notice_processor do |message|
      say message
    end
    conn
  end
rescue PG::ConnectionBad => e
  abort e.message
rescue URI::InvalidURIError
  abort "Invalid url"
end
create_table(name) click to toggle source
# File lib/pgslice/helpers.rb, line 163
def create_table(name)
  if name.include?(".")
    schema, name = name.split(".", 2)
  else
    schema = self.schema
  end
  Table.new(schema, name)
end
execute(query, params = []) click to toggle source
# File lib/pgslice/helpers.rb, line 56
def execute(query, params = [])
  connection.exec_params(query, params).to_a
end
log(message = nil) click to toggle source

output

# File lib/pgslice/helpers.rb, line 13
def log(message = nil)
  error message
end
log_sql(message = nil) click to toggle source
# File lib/pgslice/helpers.rb, line 17
def log_sql(message = nil)
  say message
end
make_fk_def(fk_def, table) click to toggle source
# File lib/pgslice/helpers.rb, line 176
def make_fk_def(fk_def, table)
  "ALTER TABLE #{quote_table(table)} ADD #{fk_def};"
end
make_index_def(index_def, table) click to toggle source
# File lib/pgslice/helpers.rb, line 172
def make_index_def(index_def, table)
  index_def.sub(/ ON \S+ USING /, " ON #{quote_table(table)} USING ").sub(/ INDEX .+ ON /, " INDEX ON ") + ";"
end
name_format(period) click to toggle source
# File lib/pgslice/helpers.rb, line 104
def name_format(period)
  case period.to_sym
  when :day
    "%Y%m%d"
  when :month
    "%Y%m"
  else
    "%Y"
  end
end
partition_date(partition, name_format) click to toggle source
# File lib/pgslice/helpers.rb, line 115
def partition_date(partition, name_format)
  DateTime.strptime(partition.name.split("_").last, name_format)
end
quote_ident(value) click to toggle source
# File lib/pgslice/helpers.rb, line 151
def quote_ident(value)
  PG::Connection.quote_ident(value)
end
quote_no_schema(table) click to toggle source
# File lib/pgslice/helpers.rb, line 159
def quote_no_schema(table)
  quote_ident(table.name)
end
quote_table(table) click to toggle source
# File lib/pgslice/helpers.rb, line 155
def quote_table(table)
  table.quote_table
end
round_date(date, period) click to toggle source
# File lib/pgslice/helpers.rb, line 119
def round_date(date, period)
  date = date.to_date
  case period.to_sym
  when :day
    date
  when :month
    Date.new(date.year, date.month)
  else
    Date.new(date.year)
  end
end
run_queries(queries) click to toggle source
# File lib/pgslice/helpers.rb, line 60
def run_queries(queries)
  connection.transaction do
    execute("SET LOCAL client_min_messages TO warning") unless options[:dry_run]
    log_sql "BEGIN;"
    log_sql
    run_queries_without_transaction(queries)
    log_sql "COMMIT;"
  end
end
run_queries_without_transaction(queries) click to toggle source
# File lib/pgslice/helpers.rb, line 82
def run_queries_without_transaction(queries)
  queries.each do |query|
    run_query(query)
  end
end
run_query(query) click to toggle source
# File lib/pgslice/helpers.rb, line 70
def run_query(query)
  log_sql query
  unless options[:dry_run]
    begin
      execute(query)
    rescue PG::ServerError => e
      abort("#{e.class.name}: #{e.message}")
    end
  end
  log_sql
end
schema() click to toggle source
# File lib/pgslice/helpers.rb, line 51
def schema
  connection # ensure called first
  @schema
end
server_version_num() click to toggle source
# File lib/pgslice/helpers.rb, line 88
def server_version_num
  execute("SHOW server_version_num")[0]["server_version_num"].to_i
end
sql_date(time, cast, add_cast = true) click to toggle source

helpers

# File lib/pgslice/helpers.rb, line 94
def sql_date(time, cast, add_cast = true)
  if cast == "timestamptz"
    fmt = "%Y-%m-%d %H:%M:%S UTC"
  else
    fmt = "%Y-%m-%d"
  end
  str = "'#{time.strftime(fmt)}'"
  add_cast ? "#{str}::#{cast}" : str
end