Script to kill old connections(sessions) on Redhsift

kill_old_redshift_sessions.rb

require 'pg'
require 'socket'
require 'pp'

TIMEOUT = 5
IDLE_CONN_TIMEOUT_MINUTES = 30 # 30 minutes
#IDLE_CONN_TIMEOUT_MINUTES = 1 # 1 minute for testing
CURRENT_SESSIONS_QUERY = "select * from stv_sessions where user_name = '%{user}' and starttime < CURRENT_TIMESTAMP - INTERVAL '%{timeout} minutes' order by starttime;"
KILL_CONNECTION_QUERY = "select pg_terminate_backend(%{pid})"

@db_conf = {
  host: '',
  port: 5439,
  user: '',
  password: '',
  dbname: '',
}

def create_connection(db_conf)
  hostaddr = IPSocket.getaddress(db_conf[:host])
  db_conf[:hostaddr] = hostaddr
  conn = PG::Connection.connect_start(db_conf)
  raise ("Unable to create a new connection.") unless conn
  raise ("Connection failed: %s" % [ conn.error_message ]) if conn.status == PG::CONNECTION_BAD

  socket = conn.socket_io
  poll_status = PG::PGRES_POLLING_WRITING
  until poll_status == PG::PGRES_POLLING_OK || poll_status == PG::PGRES_POLLING_FAILED
    case poll_status
    when PG::PGRES_POLLING_READING
      IO.select([socket], nil, nil, TIMEOUT) or raise ("Asynchronous connection timed out!(READING)")
    when PG::PGRES_POLLING_WRITING
      IO.select(nil, [socket], nil, TIMEOUT) or raise ("Asynchronous connection timed out!(WRITING)")
    end
    poll_status = conn.connect_poll
  end
  unless conn.status == PG::CONNECTION_OK
    raise ("Connect failed: %s" % [ conn.error_message.to_s.lines.uniq.join(" ") ])
  end
  conn
end

def send_query(sql, options = {})
  conn = options[:conn] || create_connection(@db_conf)
  conn.exec(sql)
ensure
  conn.finish if conn && options[:conn].nil?
end

def get_old_and_idle_processes(conn)
  sql = CURRENT_SESSIONS_QUERY % {user: @db_conf[:user], timeout: IDLE_CONN_TIMEOUT_MINUTES}
  conn.exec(sql).collect do |row|
    row['process'].to_s.strip.to_i
  end
end

def kill_process(conn, pid)
  sql = KILL_CONNECTION_QUERY % {pid: pid}
  conn.exec(sql)
end

def check_and_kill_connections
  conn = create_connection(@db_conf)
  processes = get_old_and_idle_processes(conn)
  processes.each do |pid|
    puts "Killing process - pid:#{pid}"
    kill_process(conn, pid)
  end
ensure
  conn.finish if conn
end

check_and_kill_connections