SampleListener.java

package net.ttddyy.dsproxy.example;

import net.ttddyy.dsproxy.ExecutionInfo;
import net.ttddyy.dsproxy.QueryInfo;
import net.ttddyy.dsproxy.listener.QueryExecutionListener;

import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * Original here: https://corner.squareup.com/2016/01/query-sniper.html
 *
 * @author Tadaya Tsuyukubo
 */
public class SampleListener implements QueryExecutionListener {
    private static final long MAX_QUERY_TIME = 10_000; // 10 seconds
    private final ScheduledExecutorService scheduledExecutorService;

    public SampleListener(ScheduledExecutorService scheduledExecutorService) {
        this.scheduledExecutorService = scheduledExecutorService;
    }

    @Override
    public void beforeQuery(ExecutionInfo execInfo, List<QueryInfo> queryInfoList) {
        long timeBudget = MAX_QUERY_TIME;

        Statement statement = execInfo.getStatement();

        // TODO: need some way to toggle queryRunning flag when it's done. such as threadlocal, map, etc.
        // currently, same instance of QueryExecutionListener is used for every jdbc invocation.
        // thus cannot simply have instance variable to pass the query running info from beforeQuery to afterQuery
        boolean queryRunning = true;

        scheduledExecutorService.schedule(() -> {
            if (queryRunning) {
                try {
                    statement.cancel();
                } catch (SQLException e) {
                    // TODO
                }
//                queryRunning = false;
            }
        }, timeBudget, TimeUnit.MILLISECONDS);
    }


    @Override
    public void afterQuery(ExecutionInfo execInfo, List<QueryInfo> queryInfoList) {
        // TODO: make the query running flag set to false
    }
}