net.snowflake.client.HeartbeatIT Maven / Gradle / Ivy

Go to download
/*
 * Test heartbeat
 */
package net.snowflake.client;

import net.snowflake.client.jdbc.SnowflakeSQLException;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.logging.Level;
import java.util.logging.Logger;

import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;

import static org.junit.Assert.*;


/**
 * This test assumes that GS has been set up
 *
 * @author jhuang
 */
public class HeartbeatIT extends AbstractDriverIT
{
  private static Logger logger =
  Logger.getLogger(HeartbeatIT.class.getName());

  public HeartbeatIT()
  {
  }

  /**
   * set up:
   *
   * change the master token validity to 10 seconds
   * change the session token validity to 5 seconds
   * change the SESSION_RECORD_ACCESS_INTERVAL_SECS to 1 second
   *
   * @throws Exception
   */
  @BeforeClass
  public static void setUpClass() throws Exception
  {
    outputLine("[ HeartbeatIT ] start suite");
    outputLine("HeartbeatIT : running");
    outputLine("running driver HeartbeatIT");

    try
    {
      Class.forName("net.snowflake.client.jdbc.SnowflakeDriver");
    }
    catch (ClassNotFoundException ex)
    {
      logger.log(Level.SEVERE, "Driver not found", ex);
    }

    try
    {
      Connection connection = getSnowflakeConnection();

      Statement statement = connection.createStatement();
      statement.executeUpdate("alter system set master_token_validity=60");
      statement.close();

      statement.executeUpdate("alter system set session_token_validity=20");
      statement.close();

      statement.executeUpdate("alter system set SESSION_RECORD_ACCESS_INTERVAL_SECS=1");
      statement.close();

      connection.close();
    }
    catch (Throwable ex)
    {
      ex.printStackTrace();
      throw ex;
    }
  }

  /**
   * Reset master_token_validity, session_token_validity,
   * SESSION_RECORD_ACCESS_INTERVAL_SECS to default.
   *
   * @throws SQLException
   */
  @AfterClass
  public static void tearDownClass() throws SQLException
  {
    outputLine("[ HeartbeatIT ] done suite " +
               diffs.toString() + " diff " +
               succs.toString() + " suc");

    try
    {
      Connection connection = getSnowflakeConnection();

      Statement statement = connection.createStatement();
      statement.executeUpdate("alter system set master_token_validity=default");
      statement.close();

      statement.executeUpdate("alter system set session_token_validity=default");
      statement.close();

      statement.executeUpdate(
              "alter system set SESSION_RECORD_ACCESS_INTERVAL_SECS=default");
      statement.close();

      connection.close();
    }
    catch (Throwable ex)
    {
      ex.printStackTrace();
      throw ex;
    }
  }

  /**
   * Get connection using snowflake account
   *
   * @return
   * @throws SQLException
   */
  private static Connection getSnowflakeConnection()
          throws SQLException
  {
    String accountStr = "snowflake";

    // build connection properties
    Properties properties = new Properties();

    properties.put("user", "admin");
    properties.put("password", "test");
    properties.put("account", accountStr);
    properties.put("ssl", "off");
    properties.put("CLIENT_SESSION_KEEP_ALIVE", "true");
    properties.put("tracing", "finest");

    properties.put("useV1QueryAPI", "on");

    // create a new connection
    String connectStr = System.getenv("SF_JDBC_CONNECT_STRING");

    // use the default connection string if it is not set in environment
    if(connectStr == null)
    {
      connectStr = "jdbc:snowflake://localhost:8080";
    }

    return DriverManager.getConnection(connectStr, properties);
  }

  /**
   * create a new connection with or without keep alive session and submit
   * a query that will take longer than the master token validity.
   *
   * @param useKeepAliveSession
   * @return
   * @throws SQLException
   */
  private Boolean submitQuery(boolean useKeepAliveSession,
                              int queryIdx) throws SQLException,
                                                   InterruptedException
  {
    Connection connection = null;
    Statement statement = null;
    ResultSet resultSet = null;
    ResultSetMetaData resultSetMetaData = null;

    try
    {
      Properties sessionParams = new Properties();

      if (useKeepAliveSession)
      {
        sessionParams.put("CLIENT_SESSION_KEEP_ALIVE", "true");
      }

      connection = getConnection(GET_ACCOUNT_FROM_ENV,
                                         sessionParams);

      statement = connection.createStatement();

      // sleep 61 seconds
      Thread.sleep(61000);

      // execute query
      resultSet = statement.executeQuery("SELECT 1");

      resultSetMetaData = resultSet.getMetaData();

      // assert column count
      assertTrue(resultSetMetaData.getColumnCount() == 1);

      // assert we get 1 row
      assertTrue(resultSet.next());

      logger.info("Query " + queryIdx + " passed ");

      statement.close();
    }
    finally
    {
      closeSQLObjects(resultSet, statement, connection);
    }
    return true;
  }

  /**
   * Test heartbeat by starting 10 threads. Each get a connection and wait
   * for a time longer than master token validity and issue a query to make
   * sure the query succeeds.
   *
   * @throws SQLException
   */
  @Test
  public void testSuccess() throws SQLException
  {
    outputLine("running testSuccess");

    int concurrency = 5;

    ExecutorService executorService = Executors.newFixedThreadPool(10);

    List futures = new ArrayList();

    try
    {
      // create 10 threads, each open a connection and submit a query
      // after sleeping 15 seconds
      for(int idx = 0; idx < concurrency; idx++)
      {
        logger.info("open a new connection and submit query " + idx);

        final int queryIdx = idx;

        futures.add(executorService.submit(new Callable()
        {
          @Override
          public Boolean call() throws Exception
          {
            return submitQuery(true, queryIdx);
          }
        }
        ));
      };

      executorService.shutdown();

      for (int idx = 0; idx < concurrency; idx++)
        futures.get(idx).get();

      succs.incrementAndGet();
      outputLine("testSuccess: success");
    }
    catch (Throwable ex)
    {
      diffs.incrementAndGet();
      outputLine("testSuccess: difference");

      logger.log(Level.SEVERE, "Test failed with exception: ", ex);
      assertTrue(ex.getLocalizedMessage(), false);
    }
  }

  /**
   * Test no heartbeat by starting 1 thread. It gets a connection and wait
   * for a time longer than master token validity and issue a query to make
   * sure the query fails.
   *
   * @throws SQLException
   */
  @Test
  public void testFailure() throws SQLException
  {
    outputLine("running testFailure");

    ExecutorService executorService = Executors.newFixedThreadPool(10);

    List futures = new ArrayList();

    try
    {
      logger.info("open a new connection and submit query");

      futures.add(executorService.submit(new Callable()
      {
        @Override
        public Boolean call() throws Exception
        {
          return submitQuery(false, 0);
        }
      }
      ));

      executorService.shutdown();

      futures.get(0).get();

      diffs.incrementAndGet();
      outputLine("testFailure: difference");
    }
    catch (Throwable ex)
    {
      Throwable rootCause = ex;

      while (rootCause.getCause() != null)
        rootCause = rootCause.getCause();

      if ((rootCause instanceof SnowflakeSQLException) &&
          ((SnowflakeSQLException)rootCause).getErrorCode() == 390114)
      {
        succs.incrementAndGet();
        outputLine("testFailure: success");
      }
      else
      {
        diffs.incrementAndGet();
        outputLine("testFailure: difference");

        logger.log(Level.SEVERE, "Test failed with exception: ", ex);
        assertTrue(ex.getLocalizedMessage(), false);
      }
    }
  }

}