Killing Tez Submitted Queries

ClosingPoolableConnectionsDataSource.java

public class ClosingPoolableConnectionsDataSource extends BasicDataSource implements ConnectionRegister {

  private Set<PoolableConnection> connections = new HashSet<>();

  @Override
  public synchronized void register(Connection conn) {
    if (PoolableConnection.class.isAssignableFrom(conn.getClass())) {
      connections.add((PoolableConnection) conn);
    }
  }

  @Override
  public synchronized void unregister(Connection conn) {
    if (PoolableConnection.class.isAssignableFrom(conn.getClass())) {
      connections.remove(conn);
    }
  }

  @Override
  public synchronized void close() throws java.sql.SQLException {
    for (PoolableConnection conn : connections) {
      conn.reallyClose();
    }
    connections.clear();
    super.close();
  }
}