/*
 * Copyright 2014-2019 JKOOL, LLC.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.jkoolcloud.client.api.utils;

import java.util.concurrent.TimeUnit;

import com.jkoolcloud.client.api.service.JKQueryAsync;
import com.jkoolcloud.client.api.service.JKQueryHandle;
import com.jkoolcloud.client.api.service.JKRetryConnectionHandler;
import com.jkoolcloud.client.api.service.JKTraceConnectionHandler;
import com.jkoolcloud.client.api.service.JKTraceQueryCallback;

public class JKQLCmd {
	public static void main(String[] args) {
		JKCmdOptions options = new JKCmdOptions(JKQLCmd.class, args);
		if (options.usage != null) {
			System.out.println(options.usage);
			System.exit(-1);
		}
		options.print();
		JKTraceQueryCallback callback = new JKTraceQueryCallback(System.out, options.json_path, options.trace);
		try (JKQueryAsync jkQueryAsync = new JKQueryAsync(System.getProperty("jk.ws.uri", options.uri), System.getProperty("jk.access.token", options.token))) 
		{
			// setup jKool WebSocket connection and connect
			if (options.retryTimeMs > 0) {
				jkQueryAsync.addConnectionHandler(new JKRetryConnectionHandler(options.retryTimeMs, TimeUnit.MILLISECONDS));
			}
			jkQueryAsync.setTimeZone(options.timezone)
				.setDateFilter(options.daterange)
				.setRepoId(options.reponame)
				.setTrace(options.trace);

			jkQueryAsync.addConnectionHandler(new JKTraceConnectionHandler(System.out, options.trace));
			jkQueryAsync.addDefaultCallbackHandler(callback);
			jkQueryAsync.connect();
			

			// run query in async mode with a callback
			JKQueryHandle qhandle = jkQueryAsync.callAsync(options.query, options.maxRows, callback);
			System.out.println("Running query=[" + qhandle.getQuery() + "]");
			System.out.println("Prepared statement=[" + qhandle.getStatement() + "]");
			if (!qhandle.isSubscribeQuery()) {
				// standard query only one response expected
				qhandle.awaitOnDone(options.waitTimeMs, TimeUnit.MILLISECONDS);
			} else {
				// streaming query, so lets collect responses until timeout
				Thread.sleep(options.waitTimeMs);
				System.out.println("Cancelling query=[" + qhandle.getQuery() + "], id=" + qhandle.getId() + ", mid=" + qhandle.getLastMsgId());
				qhandle = jkQueryAsync.cancelAsync(qhandle);
				if (qhandle != null) {
					qhandle.awaitOnDone(options.waitTimeMs, TimeUnit.MILLISECONDS);
				}
			}
		} catch (Throwable e) {
			System.err.println("Failed to execute: " + options.toString());
			e.printStackTrace();
		} finally {
			System.out.println("Stats: msg.recvd=" + callback.getMsgCount() + ", err.count=" + callback.getErrorCount());
		} 
	}
}

@@ -635,16 +635,44 @@
 utils;%0A%0A
+import java.io.IOException;%0A
 import j
@@ -701,16 +701,16 @@
 meUnit;%0A
-
 %0Aimport 
@@ -2186,16 +2186,78 @@
 lback);%0A
+%09%09%09Runtime.getRuntime().addShutdownHook(new VMStop(qhandle));%0A
 %09%09%09Syste
@@ -3196,16 +3196,16 @@
 unt());%0A
-
 %09%09%7D %0A%09%7D%0A
@@ -3206,8 +3206,559 @@
 %7D %0A%09%7D%0A%7D%0A
+%0Aclass VMStop extends Thread %7B%0A%09JKQueryHandle handle;%0A%09%0A%09protected VMStop(JKQueryHandle h) %7B%0A%09%09handle = h;%0A%09%7D%0A%09%0A%09@Override%0A%09public void run() %7B%0A%09%09try %7B%0A%09%09%09if (handle.isSubscribeQuery()) %7B%0A%09%09%09%09System.out.println(%22VM Stopping. Cancel subscription on handle= %22 + handle);%0A%09%09%09%09handle.cancelAsync();%0A%09%09%09%09handle.awaitOnDone(5000, TimeUnit.MILLISECONDS);%0A%09%09%09%7D%0A%09%09%09System.out.println(%22VM Stopping. Closing handle=%22 + handle);%0A%09%09%09handle.getStatement().getQueryAsync().close();%0A%09%09%7D catch (IOException %7C InterruptedException e) %7B%0A%09%09%09e.printStackTrace();%0A%09%09%7D%0A%09%7D%0A%7D%0A

handle VM shutdown example