/*
    Program which times SQL communiation latency by 
    firing off a number of simple SQL statements which don't
    perform I/O.

    Versioning: $Id: LatencyTester.java 419 2010-05-21 22:46:23Z troels $
*/
import java.sql.*;

import java.io.*;
import java.util.Random;

public class LatencyTester {

    // some "global" variables, to save code
    static String dbname;
    static String servername;
    static String username;
    static String sqlprofile;
    static int nStatements = 10000;

    public static void usage() {
        echo("Usage:");
        echo("LatencyTester <type> <dbname> [ <servername> <username> ]");
        echo("");
        echo("The type parameter may be one of:");
        echo(" short - short pieces of data are being used");
        echo(" long  - long  pieces of data are being used");
    }

    public static void err(String message) {
        System.err.println("Error: "+message);
        System.err.println("Provide the -h argument for help");
        System.exit(1);
    }

    // To save typing:
    public static void echo(String message) {
        System.out.println(message);
    }

    public static Connection getConn()
        throws
            SQLException,
            ClassNotFoundException,
            InstantiationException,
            IllegalAccessException
    {
        Connection conn = null;

        if (null == servername) {
            echo("Using local IPC connection");

            Class.forName("COM.ibm.db2.jdbc.app.DB2Driver").newInstance();
            conn = DriverManager.getConnection("jdbc:db2:"+dbname);
        } else {
            echo("Using TCP connection");

            Console cons = System.console();
            echo("Please enter password:");
            String password = new String(cons.readPassword());

            Class.forName("com.ibm.db2.jcc.DB2Driver").newInstance();
            String options = ":user="+username+";password="+password+";";
            String url = "jdbc:db2://"+servername+":50000/"+dbname+options;
            conn = DriverManager.getConnection(url);
        }

        return conn;
    }

    public static void runQueries(Connection conn, String sqlprofile)
        throws SQLException
    {
        ResultSet resSet;

        echo("Running "+nStatements+" statements; profile='"+sqlprofile+"'");

        if (sqlprofile.equals("short")) {

            // +1 enables prepareStatement to deduce ?'s type
            String sql = "VALUES(? + 1)";
            PreparedStatement pstmt = conn.prepareStatement(sql);

            // Shouldn't be blocked by lack of entropy; verified
            // by noticing that strace doesn't open /dev/random
            // (in contrast to what SecureRandom seems to do).
            Random rand = new Random();

            int inVal;

            for (int i=0; i<nStatements; i++) {
                inVal = rand.nextInt();
                pstmt.setInt(1,inVal);
                resSet = pstmt.executeQuery();
                if (!resSet.next()) {
                    err("Couldn't grab value from resultset");
                }
            }

        } else {

            Statement stmt = conn.createStatement();

            for (int i=0; i<nStatements; i++) {
                resSet = stmt.executeQuery(buildLongSQL());
                if (!resSet.next()) {
                    err("Couldn't grab value from resultset");
                }
            }

        }
    }

    // Returns a random 2000 char string
    public static String buildLongSQL() {
        // Build a string from a random selection of sub-strings
        String[] strings = {
            "lksjflkvjw",
            "pvoiwepwvk",
            "voiwuel2vn",
            "v0s9dfohjw",
            "289vsod2,c"
        };
        StringBuilder sb = new StringBuilder();
        Random rand = new Random();

        sb.append("VALUES('");
        for (int i=0; i<200; i++) {
            sb.append(strings[rand.nextInt(5)]);
        }
        sb.append("')");

        return sb.toString();
    }

    // args handling and sanity checks
    public static void main (String[] args) {

        int nArgs = args.length;

        if (nArgs != 1 && nArgs != 2 && nArgs != 4) {
            err("Wrong number of arguments");
            System.exit(1);
        }

        sqlprofile = args[0];
        if (sqlprofile.equals("-h") || sqlprofile.equals("--help")) {
            usage();
            System.exit(0);
        }
        if (!(sqlprofile.equals("short") || sqlprofile.equals("long"))) {
            err("Unknown test profile: " + sqlprofile);
        }
        if (nArgs == 1) {
            err("Wrong number of arguments");
            System.exit(1);
        }

        dbname = args[1];

        if (nArgs == 4) {
            servername = args[2];
            username = args[3];
        }

        long start = 0;
        long end   = 0;
        try {
            Connection conn = getConn();
            start = System.currentTimeMillis();
            runQueries(conn, sqlprofile);
            end = System.currentTimeMillis();
            conn.close();
        } catch (Exception e) {
            err("Got exception: "+e);
        }

        long diff = end - start;
        echo("Number of ms elapsed (excl connection setup): "+diff);
    }
}