class PgConn::SessionMethods

def disable(database)
  !database.nil? or raise ArgumentError
  conn.execute "alter database #{database} allow_connections = false"
end

def disable_triggers()
  conn.execute "set session session_replication_role = replica"
end

def enable(database)
  !database.nil? or raise ArgumentError
  conn.execute "alter database #{database} allow_connections = true"
end

def enable_triggers()
  conn.execute "set session session_replication_role = DEFAULT"
end

def enabled?(database)
  !database.nil? or raise ArgumentError
  conn.value "select datallowconn from pg_catalog.pg_database where datname = '#{database}'"
end

def exclusive(database, &block)
  !database.nil? or raise ArgumentError
  begin
    disable(database)
    terminate(database, nil)
    yield
  ensure
    enable(database)
  end
end

def initialize(conn)
  @conn = conn
end

def list(database)
  if database
    conn.values "select usename from pg_stat_activity where datname = '#{database}'"
  else
    conn.tuples %(
      select datname, usename
      from pg_stat_activity
      where datname is not null and usename is not null
    )
  end
end

def pids(database, users)
  users ||= list(database)
  if !users.empty?
    users_sql = "(" + users.map { |user| "'#{user}'" }.join(", ") + ")"
    conn.values "select pid from pg_stat_activity where datname = '#{database}' and usename in #{users_sql}"
  else
    conn.values "select pid from pg_stat_activity where datname = '#{database}'"
  end
end

def terminate(database, *users)
  users = Array(users).flatten
  !database.nil? or raise ArgumentError
  enabled = self.enabled?(database)
  case users
    when [];
      return
    when [nil]
      self.disable(database) if enabled
      users = self.list(database)
    else
      users = Array(users).flatten
  end
  pids = self.pids(database, users)
  if !pids.empty?
    pids_sql = pids.map { |pid| "(#{pid})" }.join(", ")
    conn.execute "select pg_terminate_backend(pid) from ( values #{pids_sql} ) as x(pid)"
  end
  self.enable(database) if self.enabled?(database) != enabled
end

def triggers(on_off, &block)
  begin
    active = triggers?
    if on_off && !active
      enable_triggers
    elsif !on_off && active
      disable_triggers
    end
    yield
  ensure
    case active
      when true; enable_triggers if !triggers?
      when false; disable_triggers if triggers?
    end
  end
end

def triggers?() conn.value "select current_setting('session_replication_role') <> 'replica'" end