Redshift Connection Killer

redshift_connection_killer.py

#!/usr/bin/env python

import sys, os
import json
import argparse
from classes.Config import Config
from classes.pg_tools import PGInteraction
from classes.Logger import Logger as l
from pprint import pprint
from datetime import datetime, timedelta


l = l()

RESERVED_USERS = ['admin', 'alooma', 'rdsdb']

conf = Config()
pg_db_name = conf['cosmo']['db_name']
pg_user = conf['cosmo']['user']
pg_host = conf['cosmo']['host']
pg_password = conf['cosmo']['password']
pg_port = conf['cosmo']['port']
pg = PGInteraction(dbname=pg_db_name,host=pg_host,user=pg_user,password=pg_password,port=pg_port,schema='public')
pg.conn()
pg.batchOpen()

c_date = datetime.utcnow() - timedelta(minutes=60)

sql = """
    select process from stv_sessions
    where starttime <= '%(a)s'
      and user_name not in ('%(u)s') and db_name <> 'dev'
      and process  not in
        ( select pid
          from stv_recents
          where status = 'Running')
      and process not in
        ( select pid
          from stl_query
          where starttime >= '%(a)s'
            and database <>'dev')""" % ({'a': c_date, 'u': "','".join((str(n) for n in RESERVED_USERS))})

l.l('\n' + sql + '\n')

result = pg.fetch_sql(sql)

kill_list = [x[0] for x in result]

kill_len = len(kill_list)

l.l(str(kill_len) + " connections to be killed")

if kill_len == 0:
    l.l("no work to do")
    exit()

for k in kill_list:
    sql_kill =  'select pg_terminate_backend(%s);' % (str(k))
    pg.exec_sql(sql_kill)
    l.l("connection %s killed: " % str(k))

l.l('kill list complete')