/*
 * Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
 */
package net.snowflake.client.jdbc;

import net.snowflake.client.ConditionalIgnoreRule.ConditionalIgnore;
import net.snowflake.client.RunningNotOnTestaccount;
import net.snowflake.client.RunningOnGithubAction;
import net.snowflake.client.category.TestCategoryConnection;
import net.snowflake.client.core.QueryStatus;
import net.snowflake.client.jdbc.telemetryOOB.TelemetryService;
import net.snowflake.common.core.SqlState;
import org.apache.commons.codec.binary.Base64;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLClientInfoException;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.sql.Statement;
import java.util.Collections;
import java.util.Enumeration;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import static net.snowflake.client.core.QueryStatus.RUNNING;
import static net.snowflake.client.core.SessionUtil.CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY;
import static org.hamcrest.CoreMatchers.equalTo;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

/**
 * Connection integration tests
 */
@Category(TestCategoryConnection.class)
public class ConnectionIT extends BaseJDBCTest
{
  // create a local constant for this code for testing purposes (already defined in GS)
  private static final int INVALID_CONNECTION_INFO_CODE = 390100;
  private static final int SESSION_CREATION_OBJECT_DOES_NOT_EXIST_NOT_AUTHORIZED = 390201;
  private static final int ROLE_IN_CONNECT_STRING_DOES_NOT_EXIST = 390189;

  private static final int WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS = 5000;

  String errorMessage = null;

  private boolean defaultState;

  @Rule
  public TemporaryFolder tmpFolder = new TemporaryFolder();

  @Before
  public void setUp()
  {
    TelemetryService service = TelemetryService.getInstance();
    service.updateContextForIT(getConnectionParameters());
    defaultState = service.isEnabled();
    service.setNumOfRetryToTriggerTelemetry(3);
    service.enable();
  }

  @After
  public void tearDown() throws InterruptedException
  {
    TelemetryService service = TelemetryService.getInstance();
    // wait 5 seconds while the service is flushing
    TimeUnit.SECONDS.sleep(5);

    if (defaultState)
    {
      service.enable();
    }
    else
    {
      service.disable();
    }
    service.resetNumOfRetryToTriggerTelemetry();
  }

  @Test
  public void testSimpleConnection() throws SQLException
  {
    Connection con = getConnection();
    Statement statement = con.createStatement();
    ResultSet resultSet = statement.executeQuery("show parameters");
    assertTrue(resultSet.next());
    assertFalse(con.isClosed());
    statement.close();
    con.close();
    assertTrue(con.isClosed());
    con.close(); // ensure no exception
  }

  @Test
  @Ignore
  public void test300ConnectionsWithSingleClientInstance() throws SQLException
  {
    // concurrent testing
    int size = 300;
    Connection con = getConnection();
    String database = con.getCatalog();
    String schema = con.getSchema();
    con.createStatement().execute("create or replace table bigTable(rowNum number,rando " +
                                  "number) as (select seq4()," +
                                  "uniform(1, 10, random()) from table(generator(rowcount=>10000000)) v)");
    con.createStatement().execute("create or replace table conTable(colA number)");

    ExecutorService taskRunner = Executors.newFixedThreadPool(size);
    for (int i = 0; i < size; i++)
    {
      ConcurrentConnections newTask = new ConcurrentConnections();
      taskRunner.submit(newTask);
    }
    assertEquals(null, errorMessage);
    taskRunner.shutdownNow();
  }

  /**
   * Test that login timeout kick in when connect through datasource
   *
   * @throws SQLException if any SQL error occurs
   */
  @Test
  public void testLoginTimeoutViaDataSource() throws SQLException
  {
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
    ds.setUrl("jdbc:snowflake://fakeaccount.snowflakecomputing.com");
    ds.setUser("fakeUser");
    ds.setPassword("fakePassword");
    ds.setAccount("fakeAccount");
    ds.setLoginTimeout(10);

    long startLoginTime = System.currentTimeMillis();
    try
    {
      ds.getConnection();
      fail();
    }
    catch (SQLException e)
    {
      assertThat(e.getErrorCode(), is(ErrorCode.NETWORK_ERROR.getMessageCode()));
    }
    long endLoginTime = System.currentTimeMillis();

    assertTrue(endLoginTime - startLoginTime < 30000);
  }

  /**
   * Test production connectivity in case cipher suites or tls protocol change
   * Use fake username and password but correct url
   * Expectation is receiving incorrect username or password response from server
   */
  @Test
  public void testProdConnectivity() throws SQLException
  {
    String[] deploymentUrls = {
        "jdbc:snowflake://sfcsupport.snowflakecomputing.com",
        "jdbc:snowflake://sfcsupportva.us-east-1.snowflakecomputing.com",
        "jdbc:snowflake://sfcsupporteu.eu-central-1.snowflakecomputing.com"};

    Properties properties = new Properties();

    properties.put("user", "fakesuer");
    properties.put("password", "fakepwd");
    properties.put("account", "fakeaccount");

    for (String url : deploymentUrls)
    {
      try
      {
        DriverManager.getConnection(url, properties);
        fail();
      }
      catch (SQLException e)
      {
        assertThat(e.getErrorCode(), is(INVALID_CONNECTION_INFO_CODE));
      }
    }
  }

  @Test
  public void testSetCatalogSchema() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      String db = connection.getCatalog();
      String schema = connection.getSchema();
      connection.setCatalog(db);
      connection.setSchema("PUBLIC");

      // get the current schema
      ResultSet rst = connection.createStatement().executeQuery("select current_schema()");
      assertTrue(rst.next());
      assertEquals("PUBLIC", rst.getString(1));
      assertEquals(db, connection.getCatalog());
      assertEquals("PUBLIC", connection.getSchema());

      // get the current schema
      connection.setSchema(schema);
      rst = connection.createStatement().executeQuery("select current_schema()");
      assertTrue(rst.next());
      assertEquals(schema, rst.getString(1));
      rst.close();
    }
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testConnectionGetAndSetDBAndSchema() throws SQLException
  {
    Connection con = getConnection();

    final String database = System.getenv("SNOWFLAKE_TEST_DATABASE").toUpperCase();
    final String schema = System.getenv("SNOWFLAKE_TEST_SCHEMA").toUpperCase();

    assertEquals(database, con.getCatalog());
    assertEquals(schema, con.getSchema());

    final String SECOND_DATABASE = "SECOND_DATABASE";
    final String SECOND_SCHEMA = "SECOND_SCHEMA";
    Statement statement = con.createStatement();
    statement.execute(String.format("create or replace database %s", SECOND_DATABASE));
    statement.execute(String.format("create or replace schema %s", SECOND_SCHEMA));
    statement.execute(String.format("use database %s", database));

    // TODO: use the other database and schema
    con.setCatalog(SECOND_DATABASE);
    assertEquals(SECOND_DATABASE, con.getCatalog());
    assertEquals("PUBLIC", con.getSchema());

    con.setSchema(SECOND_SCHEMA);
    assertEquals(SECOND_SCHEMA, con.getSchema());

    statement.execute(String.format("use database %s", database));
    statement.execute(String.format("use schema %s", schema));

    assertEquals(database, con.getCatalog());
    assertEquals(schema, con.getSchema());

    statement.execute(String.format("drop database if exists %s", SECOND_DATABASE));
    con.close();
  }

  @Test
  public void testConnectionClientInfo() throws SQLException
  {
    try (Connection con = getConnection())
    {
      Properties property = con.getClientInfo();
      assertEquals(0, property.size());
      Properties clientInfo = new Properties();
      clientInfo.setProperty("name", "Peter");
      clientInfo.setProperty("description", "SNOWFLAKE JDBC");
      try
      {
        con.setClientInfo(clientInfo);
        fail("setClientInfo should fail for any parameter.");
      }
      catch (SQLClientInfoException e)
      {
        assertEquals(SqlState.INVALID_PARAMETER_VALUE, e.getSQLState());
        assertEquals(200047, e.getErrorCode());
        assertEquals(2, e.getFailedProperties().size());
      }
      try
      {
        con.setClientInfo("ApplicationName", "valueA");
        fail("setClientInfo should fail for any parameter.");
      }
      catch (SQLClientInfoException e)
      {
        assertEquals(SqlState.INVALID_PARAMETER_VALUE, e.getSQLState());
        assertEquals(200047, e.getErrorCode());
        assertEquals(1, e.getFailedProperties().size());
      }
    }
  }

  // only support get and set
  @Test
  public void testNetworkTimeout() throws SQLException
  {
    Connection con = getConnection();
    int millis = con.getNetworkTimeout();
    assertEquals(0, millis);
    con.setNetworkTimeout(null, 200);
    assertEquals(200, con.getNetworkTimeout());
    con.close();
  }

  @Test
  public void testAbort() throws SQLException
  {
    Connection con = getConnection();
    assertTrue(!con.isClosed());
    con.abort(null);
    assertTrue(con.isClosed());
  }

  @Test
  public void testSetQueryTimeoutInConnectionStr() throws SQLException
  {
    Properties properties = new Properties();
    properties.put("queryTimeout", "5");
    Connection connection = getConnection(properties);
    Statement statement = connection.createStatement();
    try
    {
      statement.executeQuery("select count(*) from table(generator(timeLimit => 1000000))");
    }
    catch (SQLException e)
    {
      assertTrue(true);
      assertEquals(SqlState.QUERY_CANCELED, e.getSQLState());
      assertEquals("SQL execution canceled", e.getMessage());
    }
    statement.close();
    connection.close();
  }

  @Test
  public void testHttpsLoginTimeoutWithSSL() throws InterruptedException
  {
    long connStart = 0, conEnd;
    Properties properties = new Properties();
    properties.put("account", "wrongaccount");
    properties.put("loginTimeout", "5");
    properties.put("user", "fakeuser");
    properties.put("password", "fakepassword");
    // only when ssl is on can trigger the login timeout
    // ssl is off will trigger 404
    properties.put("ssl", "on");
    try
    {
      connStart = System.currentTimeMillis();
      Map<String, String> params = getConnectionParameters();
      // use wrongaccount in url
      String host = params.get("host");
      String[] hostItems = host.split("\\.");
      String wrongUri =
          params.get("uri").replace("://" + hostItems[0], "://wrongaccount");
      DriverManager.getConnection(wrongUri, properties);
    }
    catch (SQLException e)
    {
      assertThat("Communication error", e.getErrorCode(),
                 equalTo(ErrorCode.NETWORK_ERROR.getMessageCode()));

      conEnd = System.currentTimeMillis();
      assertThat("Login time out not taking effective",
                 conEnd - connStart < 60000);
      Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS);
      if (TelemetryService.getInstance().isDeploymentEnabled())
      {
        assertThat("Telemetry event has not been reported successfully. Error: " +
                   TelemetryService.getInstance().getLastClientError(),
                   TelemetryService.getInstance().getClientFailureCount(), equalTo(0));
      }
      return;
    }
    fail();
  }

  @Test
  public void testHttpsLoginTimeoutWithOutSSL() throws InterruptedException
  {
    Properties properties = new Properties();
    properties.put("account", "wrongaccount");
    properties.put("loginTimeout", "20");
    properties.put("user", "fakeuser");
    properties.put("password", "fakepassword");
    properties.put("ssl", "off");
    int count = TelemetryService.getInstance().getEventCount();
    try
    {
      Map<String, String> params = getConnectionParameters();
      // use wrongaccount in url
      String host = params.get("host");
      String[] hostItems = host.split("\\.");
      String wrongUri =
          params.get("uri").replace("://" + hostItems[0], "://wrongaccount");

      DriverManager.getConnection(wrongUri, properties);
    }
    catch (SQLException e)
    {
      if (TelemetryService.getInstance().getServerDeploymentName().equals(
          TelemetryService.TELEMETRY_SERVER_DEPLOYMENT.DEV.getName()) ||
          TelemetryService.getInstance().getServerDeploymentName().equals(
              TelemetryService.TELEMETRY_SERVER_DEPLOYMENT.REG.getName()))
      {
        // a connection error response (wrong user and password)
        // with status code 200 is returned in RT
        assertThat("Communication error", e.getErrorCode(),
                   equalTo(INVALID_CONNECTION_INFO_CODE));

        // since it returns normal response,
        // the telemetry does not create new event
        Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS);
        if (TelemetryService.getInstance().isDeploymentEnabled())
        {
          assertThat("Telemetry should not create new event",
                     TelemetryService.getInstance().getEventCount(), equalTo(count));
        }
      }
      else
      {
        // in qa1 and others, 404 http status code should be returned
        assertThat("Communication error", e.getErrorCode(),
                   equalTo(ErrorCode.NETWORK_ERROR.getMessageCode()));

        if (TelemetryService.getInstance().isDeploymentEnabled())
        {
          assertThat("Telemetry event has not been reported successfully. Error: " +
                     TelemetryService.getInstance().getLastClientError(),
                     TelemetryService.getInstance().getClientFailureCount(), equalTo(0));
        }
      }
      return;
    }
    fail();
  }

  @Test
  public void testWrongHostNameTimeout() throws InterruptedException
  {
    long connStart = 0, conEnd;
    Properties properties = new Properties();
    properties.put("account", "testaccount");
    properties.put("loginTimeout", "20");
    properties.put("user", "fakeuser");
    properties.put("password", "fakepassword");
    try
    {
      connStart = System.currentTimeMillis();
      Map<String, String> params = getConnectionParameters();
      // use wrongaccount in url
      String host = params.get("host");
      String[] hostItems = host.split("\\.");
      String wrongUri = params.get("uri").replace(
          "." + hostItems[hostItems.length - 2] + ".", ".wronghostname.");

      DriverManager.getConnection(wrongUri, properties);
    }
    catch (SQLException e)
    {
      assertThat("Communication error", e.getErrorCode(),
                 equalTo(ErrorCode.NETWORK_ERROR.getMessageCode()));

      conEnd = System.currentTimeMillis();
      assertThat("Login time out not taking effective",
                 conEnd - connStart < 60000);

      Thread.sleep(WAIT_FOR_TELEMETRY_REPORT_IN_MILLISECS);
      if (TelemetryService.getInstance().isDeploymentEnabled())
      {
        assertThat("Telemetry event has not been reported successfully. Error: " +
                   TelemetryService.getInstance().getLastClientError(),
                   TelemetryService.getInstance().getClientFailureCount(), equalTo(0));
      }
      return;
    }
    fail();
  }

  @Test
  public void testConnectViaDataSource() throws SQLException
  {
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();

    Map<String, String> params = getConnectionParameters();
    String account = params.get("account");
    String host = params.get("host");
    String port = params.get("port");
    String user = params.get("user");
    String password = params.get("password");
    String ssl = params.get("ssl");

    String connectStr = String.format("jdbc:snowflake://%s:%s", host, port);

    ds.setUrl(connectStr);
    ds.setAccount(account);
    ds.setSsl("on".equals(ssl));

    Connection connection = ds.getConnection(user, password);
    ResultSet resultSet = connection.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));

    connection.close();

    // get connection by server name
    // this is used by ibm cast iron studio
    ds = new SnowflakeBasicDataSource();
    ds.setServerName(params.get("host"));
    ds.setSsl("on".equals(ssl));
    ds.setAccount(account);
    ds.setPortNumber(Integer.parseInt(port));
    connection = ds.getConnection(params.get("user"), params.get("password"));
    resultSet = connection.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));

    connection.close();
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testBasicDataSourceSerialization() throws Exception
  {
    // test with username/password authentication
    // set up DataSource object and ensure connection works
    Map<String, String> params = getConnectionParameters();
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
    ds.setServerName(params.get("host"));
    ds.setSsl("on".equals(params.get("ssl")));
    ds.setAccount(params.get("account"));
    ds.setPortNumber(Integer.parseInt(params.get("port")));
    ds.setUser(params.get("user"));
    ds.setPassword(params.get("password"));
    Connection con = ds.getConnection();
    ResultSet resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    con.close();
    File serializedFile = tmpFolder.newFile("serializedStuff.ser");
    // serialize datasource object into a file
    FileOutputStream outputFile = new FileOutputStream(serializedFile);
    ObjectOutputStream out = new ObjectOutputStream(outputFile);
    out.writeObject(ds);
    out.close();
    outputFile.close();
    // deserialize into datasource object again
    FileInputStream inputFile = new FileInputStream(serializedFile);
    ObjectInputStream in = new ObjectInputStream(inputFile);
    SnowflakeBasicDataSource ds2 = (SnowflakeBasicDataSource) in.readObject();
    in.close();
    inputFile.close();
    // test connection a second time
    con = ds2.getConnection();
    resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    con.close();
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testKeyPairFileDataSourceSerialization() throws Exception
  {
    // test with key/pair authentication where key is in file
    // set up DataSource object and ensure connection works
    Map<String, String> params = getConnectionParameters();
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
    ds.setServerName(params.get("host"));
    ds.setSsl("on".equals(params.get("ssl")));
    ds.setAccount(params.get("account"));
    ds.setPortNumber(Integer.parseInt(params.get("port")));
    ds.setUser(params.get("user"));
    String privateKeyLocation = getFullPathFileInResource("encrypted_rsa_key.p8");
    ds.setPrivateKeyFile(privateKeyLocation, "test");

    // set up public key
    try (Connection con = getConnection())
    {
      Statement statement = con.createStatement();
      statement.execute("use role accountadmin");
      String pathfile = getFullPathFileInResource("encrypted_rsa_key.pub");
      String pubKey = new String(Files.readAllBytes(Paths.get(pathfile)));
      pubKey = pubKey.replace("-----BEGIN PUBLIC KEY-----", "");
      pubKey = pubKey.replace("-----END PUBLIC KEY-----", "");
      statement.execute(String.format(
          "alter user %s set rsa_public_key='%s'", params.get("user"), pubKey));
    }

    Connection con = ds.getConnection();
    ResultSet resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    con.close();
    File serializedFile = tmpFolder.newFile("serializedStuff.ser");
    // serialize datasource object into a file
    FileOutputStream outputFile = new FileOutputStream(serializedFile);
    ObjectOutputStream out = new ObjectOutputStream(outputFile);
    out.writeObject(ds);
    out.close();
    outputFile.close();
    // deserialize into datasource object again
    FileInputStream inputFile = new FileInputStream(serializedFile);
    ObjectInputStream in = new ObjectInputStream(inputFile);
    SnowflakeBasicDataSource ds2 = (SnowflakeBasicDataSource) in.readObject();
    in.close();
    inputFile.close();
    // test connection a second time
    con = ds2.getConnection();
    resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    con.close();

    // clean up
    try (Connection connection = getConnection())
    {
      Statement statement = connection.createStatement();
      statement.execute("use role accountadmin");
      statement.execute(String.format("alter user %s unset rsa_public_key",
                                      params.get("user")));
    }
  }

  @Test
  @Ignore
  public void testDataSourceOktaSerialization() throws Exception
  {
    // test with username/password authentication
    // set up DataSource object and ensure connection works
    Map<String, String> params = getConnectionParameters();
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
    ds.setServerName(params.get("host"));
    ds.setSsl("on".equals(params.get("ssl")));
    ds.setAccount(params.get("account"));
    ds.setPortNumber(Integer.parseInt(params.get("port")));
    ds.setUser(params.get("ssoUser"));
    ds.setPassword(params.get("ssoPassword"));
    ds.setAuthenticator("https://snowflakecomputing.okta.com/");
    Connection con = ds.getConnection();
    ResultSet resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    File serializedFile = tmpFolder.newFile("serializedStuff.ser");
    // serialize datasource object into a file
    FileOutputStream outputFile = new FileOutputStream(serializedFile);
    ObjectOutputStream out = new ObjectOutputStream(outputFile);
    out.writeObject(ds);
    out.close();
    outputFile.close();
    // deserialize into datasource object again
    FileInputStream inputFile = new FileInputStream(serializedFile);
    ObjectInputStream in = new ObjectInputStream(inputFile);
    SnowflakeBasicDataSource ds2 = (SnowflakeBasicDataSource) in.readObject();
    in.close();
    inputFile.close();
    // test connection a second time
    con = ds2.getConnection();
    resultSet = con.createStatement()
        .executeQuery("select 1");
    resultSet.next();
    assertThat("select 1", resultSet.getInt(1), equalTo(1));
    con.close();
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testConnectUsingKeyPair() throws Exception
  {
    Map<String, String> parameters = getConnectionParameters();
    String testUser = parameters.get("user");

    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
    SecureRandom random = SecureRandom.getInstance("SHA1PRNG");
    keyPairGenerator.initialize(2048, random);

    KeyPair keyPair = keyPairGenerator.generateKeyPair();
    PublicKey publicKey = keyPair.getPublic();
    PrivateKey privateKey = keyPair.getPrivate();

    Connection connection = getConnection();
    Statement statement = connection.createStatement();
    statement.execute("use role accountadmin");

    String encodePublicKey = Base64.encodeBase64String(publicKey.getEncoded());

    statement.execute(String.format(
        "alter user %s set rsa_public_key='%s'", testUser, encodePublicKey));

    connection.close();

    String uri = parameters.get("uri");

    Properties properties = new Properties();
    properties.put("account", parameters.get("account"));
    properties.put("user", testUser);
    properties.put("ssl", parameters.get("ssl"));
    properties.put("port", parameters.get("port"));

    // test correct private key one
    properties.put("privateKey", privateKey);
    connection = DriverManager.getConnection(uri, properties);
    connection.close();

    // test datasource connection using private key
    SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
    ds.setUrl(uri);
    ds.setAccount(parameters.get("account"));
    ds.setUser(parameters.get("user"));
    ds.setSsl("on".equals(parameters.get("ssl")));
    ds.setPortNumber(Integer.valueOf(parameters.get("port")));
    ds.setPrivateKey(privateKey);
    Connection con = ds.getConnection();
    con.close();

    // test wrong private key
    keyPair = keyPairGenerator.generateKeyPair();
    PublicKey publicKey2 = keyPair.getPublic();
    PrivateKey privateKey2 = keyPair.getPrivate();
    properties.put("privateKey", privateKey2);
    try
    {
      DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      Assert.assertEquals(390144, e.getErrorCode());
    }

    // test multiple key pair
    connection = getConnection();
    statement = connection.createStatement();
    statement.execute("use role accountadmin");

    String encodePublicKey2 = Base64.encodeBase64String(publicKey2.getEncoded());

    statement.execute(String.format(
        "alter user %s set rsa_public_key_2='%s'", testUser, encodePublicKey2));
    connection.close();

    connection = DriverManager.getConnection(uri, properties);

    // clean up
    statement = connection.createStatement();
    statement.execute("use role accountadmin");
    statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
    statement.execute(String.format("alter user %s unset rsa_public_key_2", testUser));
    connection.close();
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testPrivateKeyInConnectionString() throws SQLException, IOException
  {
    Map<String, String> parameters = getConnectionParameters();
    String testUser = parameters.get("user");

    // Test with non-password-protected private key file (.pem)
    Connection connection = getConnection();
    Statement statement = connection.createStatement();
    statement.execute("use role accountadmin");
    String pathfile = getFullPathFileInResource("rsa_key.pub");
    String pubKey = new String(Files.readAllBytes(Paths.get(pathfile)));
    pubKey = pubKey.replace("-----BEGIN PUBLIC KEY-----", "");
    pubKey = pubKey.replace("-----END PUBLIC KEY-----", "");
    statement.execute(String.format(
        "alter user %s set rsa_public_key='%s'", testUser, pubKey));
    connection.close();

    // PKCS #8
    String privateKeyLocation = getFullPathFileInResource("rsa_key.p8");
    String uri = parameters.get("uri") + "/?private_key_file=" + privateKeyLocation;
    Properties properties = new Properties();
    properties.put("account", parameters.get("account"));
    properties.put("user", testUser);
    properties.put("ssl", parameters.get("ssl"));
    properties.put("port", parameters.get("port"));
    connection = DriverManager.getConnection(uri, properties);
    connection.close();

    // PKCS #1
    privateKeyLocation = getFullPathFileInResource("rsa_key.pem");
    uri = parameters.get("uri") + "/?private_key_file=" + privateKeyLocation;
    properties = new Properties();
    properties.put("account", parameters.get("account"));
    properties.put("user", testUser);
    properties.put("ssl", parameters.get("ssl"));
    properties.put("port", parameters.get("port"));
    connection = DriverManager.getConnection(uri, properties);
    connection.close();

    // test with password-protected private key file (.p8)
    connection = getConnection();
    statement = connection.createStatement();
    statement.execute("use role accountadmin");
    pathfile = getFullPathFileInResource("encrypted_rsa_key.pub");
    pubKey = new String(Files.readAllBytes(Paths.get(pathfile)));
    pubKey = pubKey.replace("-----BEGIN PUBLIC KEY-----", "");
    pubKey = pubKey.replace("-----END PUBLIC KEY-----", "");
    statement.execute(String.format(
        "alter user %s set rsa_public_key='%s'", testUser, pubKey));
    connection.close();

    privateKeyLocation = getFullPathFileInResource("encrypted_rsa_key.p8");
    uri = parameters.get("uri") + "/?private_key_file_pwd=test&private_key_file=" + privateKeyLocation;

    connection = DriverManager.getConnection(uri, properties);
    connection.close();

    // test with incorrect password for private key
    uri = parameters.get("uri") + "/?private_key_file_pwd=wrong_password&private_key_file=" + privateKeyLocation;
    try
    {
      connection = DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertEquals((int) ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY.getMessageCode(), e.getErrorCode());
    }
    connection.close();

    // test with invalid public/private key combo (using 1st public key with 2nd private key)
    connection = getConnection();
    statement = connection.createStatement();
    statement.execute("use role accountadmin");
    pathfile = getFullPathFileInResource("rsa_key.pub");
    pubKey = new String(Files.readAllBytes(Paths.get(pathfile)));
    pubKey = pubKey.replace("-----BEGIN PUBLIC KEY-----", "");
    pubKey = pubKey.replace("-----END PUBLIC KEY-----", "");
    statement.execute(String.format(
        "alter user %s set rsa_public_key='%s'", testUser, pubKey));
    connection.close();

    privateKeyLocation = getFullPathFileInResource("encrypted_rsa_key.p8");
    uri = parameters.get("uri") + "/?private_key_file_pwd=test&private_key_file=" + privateKeyLocation;
    try
    {
      connection = DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertEquals(390144, e.getErrorCode());
    }
    connection.close();

    // test with invalid private key
    privateKeyLocation = getFullPathFileInResource("invalid_private_key.pem");
    uri = parameters.get("uri") + "/?private_key_file=" + privateKeyLocation;
    try
    {
      connection = DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertEquals((int) ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY.getMessageCode(), e.getErrorCode());
    }
    connection.close();

    //clean up
    connection = getConnection();
    statement = connection.createStatement();
    statement.execute("use role accountadmin");
    statement.execute(String.format("alter user %s unset rsa_public_key", testUser));
    connection.close();
  }

  @Test
  public void testBadPrivateKey() throws Exception
  {
    Map<String, String> parameters = getConnectionParameters();
    String testUser = parameters.get("user");

    String uri = parameters.get("uri");

    Properties properties = new Properties();
    properties.put("account", parameters.get("account"));
    properties.put("user", testUser);
    properties.put("ssl", parameters.get("ssl"));

    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("DSA");
    PrivateKey dsaPrivateKey = keyPairGenerator.generateKeyPair().getPrivate();

    try
    {
      properties.put("privateKey", "bad string");
      DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertThat(e.getErrorCode(), is(
          ErrorCode.INVALID_PARAMETER_TYPE.getMessageCode()));
    }

    try
    {
      properties.put("privateKey", dsaPrivateKey);
      DriverManager.getConnection(uri, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertThat(e.getErrorCode(), is(
          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY.getMessageCode()));
    }
  }

  @Test
  @ConditionalIgnore(condition = RunningOnGithubAction.class)
  public void testDifferentKeyLength() throws Exception
  {
    Map<String, String> parameters = getConnectionParameters();
    String testUser = parameters.get("user");

    Integer[] testCases = {2048, 4096, 8192};

    KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
    SecureRandom random = SecureRandom.getInstance("SHA1PRNG");

    for (Integer keyLength : testCases)
    {
      keyPairGenerator.initialize(keyLength, random);

      KeyPair keyPair = keyPairGenerator.generateKeyPair();
      PublicKey publicKey = keyPair.getPublic();
      PrivateKey privateKey = keyPair.getPrivate();

      Connection connection = getConnection();
      Statement statement = connection.createStatement();
      statement.execute("use role accountadmin");

      String encodePublicKey = Base64.encodeBase64String(publicKey.getEncoded());

      statement.execute(String.format(
          "alter user %s set rsa_public_key='%s'", testUser, encodePublicKey));

      connection.close();

      String uri = parameters.get("uri");

      Properties properties = new Properties();
      properties.put("account", parameters.get("account"));
      properties.put("user", testUser);
      properties.put("ssl", parameters.get("ssl"));
      properties.put("port", parameters.get("port"));
      properties.put("role", "accountadmin");

      // test correct private key one
      properties.put("privateKey", privateKey);
      connection = DriverManager.getConnection(uri, properties);

      connection.createStatement().execute(
          String.format("alter user %s unset rsa_public_key", testUser));
      connection.close();
    }
  }

  /**
   * Test production connectivity with insecure mode enabled.
   */
  @Test
  public void testInsecureMode() throws SQLException
  {
    String deploymentUrl =
        "jdbc:snowflake://sfcsupport.snowflakecomputing.com";

    Properties properties = new Properties();

    properties.put("user", "fakesuer");
    properties.put("password", "fakepwd");
    properties.put("account", "fakeaccount");
    properties.put("insecureMode", true);
    try
    {
      DriverManager.getConnection(deploymentUrl, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertThat(e.getErrorCode(), is(INVALID_CONNECTION_INFO_CODE));
    }

    deploymentUrl =
        "jdbc:snowflake://sfcsupport.snowflakecomputing.com?insecureMode=true";

    properties = new Properties();

    properties.put("user", "fakesuer");
    properties.put("password", "fakepwd");
    properties.put("account", "fakeaccount");
    try
    {
      DriverManager.getConnection(deploymentUrl, properties);
      fail();
    }
    catch (SQLException e)
    {
      assertThat(e.getErrorCode(), is(INVALID_CONNECTION_INFO_CODE));
    }
  }

  /**
   * Verify the passed memory parameters are set in the session
   */
  @Test
  public void testClientMemoryParameters() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put("CLIENT_PREFETCH_THREADS", "6");
    paramProperties.put("CLIENT_RESULT_CHUNK_SIZE", 48);
    paramProperties.put("CLIENT_MEMORY_LIMIT", 1000);
    Connection connection = getConnection(paramProperties);

    for (Enumeration<?> enums = paramProperties.propertyNames();
         enums.hasMoreElements(); )
    {
      String key = (String) enums.nextElement();
      ResultSet rs = connection.createStatement().executeQuery(
          String.format("show parameters like '%s'", key));
      rs.next();
      String value = rs.getString("value");
      assertThat(key, value, equalTo(paramProperties.get(key).toString()));
    }
  }

  /**
   * Verify the JVM memory parameters are set in the session
   */
  @Test
  public void testClientMemoryJvmParameteres() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put("CLIENT_PREFETCH_THREADS", "6");
    paramProperties.put("CLIENT_RESULT_CHUNK_SIZE", 48);
    paramProperties.put("CLIENT_MEMORY_LIMIT", 1000L);

    // set JVM parameters
    System.setProperty("net.snowflake.jdbc.clientPrefetchThreads",
                       paramProperties.get("CLIENT_PREFETCH_THREADS").toString());
    System.setProperty("net.snowflake.jdbc.clientResultChunkSize",
                       paramProperties.get("CLIENT_RESULT_CHUNK_SIZE").toString());
    System.setProperty("net.snowflake.jdbc.clientMemoryLimit",
                       paramProperties.get("CLIENT_MEMORY_LIMIT").toString());

    try
    {
      Connection connection = getConnection();

      for (Enumeration<?> enums = paramProperties.propertyNames();
           enums.hasMoreElements(); )
      {
        String key = (String) enums.nextElement();
        ResultSet rs = connection.createStatement().executeQuery(
            String.format("show parameters like '%s'", key));
        rs.next();
        String value = rs.getString("value");
        assertThat(key, value, equalTo(paramProperties.get(key).toString()));
      }
    }
    finally
    {
      System.clearProperty("net.snowflake.jdbc.clientPrefetchThreads");
      System.clearProperty("net.snowflake.jdbc.clientResultChunkSize");
      System.clearProperty("net.snowflake.jdbc.clientMemoryLimit");
    }
  }

  /**
   * Verify the connection and JVM memory parameters are set in the session.
   * The connection parameters take precedence over JVM.
   */
  @Test
  public void testClientMixedMemoryJvmParameteres() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put("CLIENT_PREFETCH_THREADS", "6");
    paramProperties.put("CLIENT_RESULT_CHUNK_SIZE", 48);
    paramProperties.put("CLIENT_MEMORY_LIMIT", 1000L);

    // set JVM parameters
    System.setProperty("net.snowflake.jdbc.clientPrefetchThreads",
                       paramProperties.get("CLIENT_PREFETCH_THREADS").toString());
    System.setProperty("net.snowflake.jdbc.clientResultChunkSize",
                       paramProperties.get("CLIENT_RESULT_CHUNK_SIZE").toString());
    System.setProperty("net.snowflake.jdbc.clientMemoryLimit",
                       paramProperties.get("CLIENT_MEMORY_LIMIT").toString());

    paramProperties.put("CLIENT_PREFETCH_THREADS", "8");
    paramProperties.put("CLIENT_RESULT_CHUNK_SIZE", 64);
    paramProperties.put("CLIENT_MEMORY_LIMIT", 2000L);

    try
    {
      Connection connection = getConnection(paramProperties);

      for (Enumeration<?> enums = paramProperties.propertyNames();
           enums.hasMoreElements(); )
      {
        String key = (String) enums.nextElement();
        ResultSet rs = connection.createStatement().executeQuery(
            String.format("show parameters like '%s'", key));
        rs.next();
        String value = rs.getString("value");
        assertThat(key, value, equalTo(paramProperties.get(key).toString()));
      }
    }
    finally
    {
      System.clearProperty("net.snowflake.jdbc.clientPrefetchThreads");
      System.clearProperty("net.snowflake.jdbc.clientResultChunkSize");
      System.clearProperty("net.snowflake.jdbc.clientMemoryLimit");
    }
  }

  /**
   * Verify the passed heartbeat frequency, which is too small, is changed to
   * the smallest valid value.
   */
  @Test
  public void testHeartbeatFrequencyTooSmall() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY, 2);
    Connection connection = getConnection(paramProperties);

    connection.getClientInfo(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY);

    for (Enumeration<?> enums = paramProperties.propertyNames();
         enums.hasMoreElements(); )
    {
      String key = (String) enums.nextElement();
      ResultSet rs = connection.createStatement().executeQuery(
          String.format("show parameters like '%s'", key));
      rs.next();
      String value = rs.getString("value");

      assertThat(key, value, equalTo("900"));
    }
  }

  /**
   * Verify the passed heartbeat frequency, which is too large, is changed to
   * the maximum valid value.
   */
  @Test
  public void testHeartbeatFrequencyTooLarge() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY, 4000);
    Connection connection = getConnection(paramProperties);

    connection.getClientInfo(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY);

    for (Enumeration<?> enums = paramProperties.propertyNames();
         enums.hasMoreElements(); )
    {
      String key = (String) enums.nextElement();
      ResultSet rs = connection.createStatement().executeQuery(
          String.format("show parameters like '%s'", key));
      rs.next();
      String value = rs.getString("value");

      assertThat(key, value, equalTo("3600"));
    }
  }

  /**
   * Verify the passed heartbeat frequency matches the output value if the
   * input is valid (between 900 and 3600).
   */
  @Test
  public void testHeartbeatFrequencyValidValue() throws Exception
  {
    Properties paramProperties = new Properties();
    paramProperties.put(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY, 1800);
    Connection connection = getConnection(paramProperties);

    connection.getClientInfo(CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY);

    for (Enumeration<?> enums = paramProperties.propertyNames();
         enums.hasMoreElements(); )
    {
      String key = (String) enums.nextElement();
      ResultSet rs = connection.createStatement().executeQuery(
          String.format("show parameters like '%s'", key));
      rs.next();
      String value = rs.getString("value");

      assertThat(key, value, equalTo(paramProperties.get(key).toString()));
    }
  }

  @Test
  public void testReadOnly() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      try
      {
        connection.setReadOnly(true);
        fail("must raise SQLFeatureNotSupportedException");
      }
      catch (SQLFeatureNotSupportedException ex)
      {
        // nop
      }

      connection.setReadOnly(false);
      connection.createStatement().execute("create or replace table readonly_test(c1 int)");
      assertFalse(connection.isReadOnly());
      connection.createStatement().execute("drop table if exists readonly_test");
    }
  }

  @Test
  public void testNativeSQL() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      // today returning the source SQL.
      assertEquals("select 1", connection.nativeSQL("select 1"));
    }
  }

  @Test
  public void testGetTypeMap() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      // return an empty type map. setTypeMap is not supported.
      assertEquals(Collections.emptyMap(), connection.getTypeMap());
    }
  }

  @Test
  public void testHolderbility() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      try
      {
        connection.setHoldability(0);
      }
      catch (SQLFeatureNotSupportedException ex)
      {
        // nop
      }
      // return an empty type map. setTypeMap is not supported.
      assertEquals(ResultSet.CLOSE_CURSORS_AT_COMMIT, connection.getHoldability());
    }
  }

  @Test
  public void testIsValid() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      assertTrue(connection.isValid(10));
      try
      {
        assertTrue(connection.isValid(-10));
        fail("must fail");
      }
      catch (SQLException ex)
      {
        // nop, no specific error code is provided.
      }
    }
  }

  @Test
  public void testUnwrapper() throws Throwable
  {
    try (Connection connection = getConnection())
    {
      boolean canUnwrap = connection.isWrapperFor(SnowflakeConnectionV1.class);
      assertTrue(canUnwrap);
      if (canUnwrap)
      {
        SnowflakeConnectionV1 sfconnection = connection.unwrap(SnowflakeConnectionV1.class);
        sfconnection.createStatement();
      }
      else
      {
        fail("should be able to unwrap");
      }
      try
      {
        connection.unwrap(SnowflakeDriver.class);
        fail("should fail to cast");
      }
      catch (SQLException ex)
      {
        // nop
      }
    }
  }

  @Test
  public void testStatementsAndResultSetsClosedByConnection() throws SQLException
  {
    Connection connection = getConnection();
    Statement statement1 = connection.createStatement();
    Statement statement2 = connection.createStatement();
    ResultSet rs1 = statement2.executeQuery("select 2;");
    ResultSet rs2 = statement2.executeQuery("select 2;");
    ResultSet rs3 = statement2.executeQuery("select 2;");
    PreparedStatement statement3 = connection.prepareStatement("select 2;");
    connection.close();
    assertTrue(statement1.isClosed());
    assertTrue(statement2.isClosed());
    assertTrue(statement3.isClosed());
    assertTrue(rs1.isClosed());
    assertTrue(rs2.isClosed());
    assertTrue(rs3.isClosed());
  }

  @Test
  public void testResultSetsClosedByStatement() throws SQLException
  {
    Connection connection = getConnection();
    Statement statement2 = connection.createStatement();
    ResultSet rs1 = statement2.executeQuery("select 2;");
    ResultSet rs2 = statement2.executeQuery("select 2;");
    ResultSet rs3 = statement2.executeQuery("select 2;");
    PreparedStatement statement3 = connection.prepareStatement("select 2;");
    ResultSet rs4 = statement3.executeQuery();
    assertFalse(rs1.isClosed());
    assertFalse(rs2.isClosed());
    assertFalse(rs3.isClosed());
    assertFalse(rs4.isClosed());
    statement2.close();
    statement3.close();
    assertTrue(rs1.isClosed());
    assertTrue(rs2.isClosed());
    assertTrue(rs3.isClosed());
    assertTrue(rs4.isClosed());
    connection.close();
  }

  @Test
  @ConditionalIgnore(condition = RunningNotOnTestaccount.class)
  public void testOKTAConnection() throws Throwable
  {
    Map<String, String> params = getConnectionParameters();
    Properties properties = new Properties();
    properties.put("user", params.get("ssoUser"));
    properties.put("password", params.get("ssoPassword"));
    properties.put("ssl", params.get("ssl"));
    properties.put("authenticator", "https://snowflakecomputing.okta.com/");

    DriverManager.getConnection(String.format(
        "jdbc:snowflake://%s.reg.snowflakecomputing.com:%s/",
        params.get("account"), params.get("port")), properties);
  }

  @Test
  @ConditionalIgnore(condition = RunningNotOnTestaccount.class)
  public void testOKTAConnectionWithOktauserParam() throws Throwable
  {
    Map<String, String> params = getConnectionParameters();
    Properties properties = new Properties();
    properties.put("user", "test");
    properties.put("password", params.get("ssoPassword"));
    properties.put("ssl", params.get("ssl"));
    properties.put("authenticator",
                   String.format("https://snowflakecomputing.okta.com;oktausername=%s;", params.get("ssoUser")));

    DriverManager.getConnection(String.format(
        "jdbc:snowflake://%s.reg.snowflakecomputing.com:%s/",
        params.get("account"), params.get("port")), properties);
  }

  @Test
  public void testValidateDefaultParameters() throws Throwable
  {
    Map<String, String> params = getConnectionParameters();
    Properties props;

    props = setCommonConnectionParameters(true);
    props.put("db", "NOT_EXISTS");
    try
    {
      DriverManager.getConnection(params.get("uri"), props);
      fail("should fail");
    }
    catch (SQLException ex)
    {
      assertEquals("error code", ex.getErrorCode(), SESSION_CREATION_OBJECT_DOES_NOT_EXIST_NOT_AUTHORIZED);
    }

    // schema is invalid
    props = setCommonConnectionParameters(true);
    props.put("schema", "NOT_EXISTS");
    try
    {
      DriverManager.getConnection(params.get("uri"), props);
      fail("should fail");
    }
    catch (SQLException ex)
    {
      assertEquals("error code", ex.getErrorCode(), SESSION_CREATION_OBJECT_DOES_NOT_EXIST_NOT_AUTHORIZED);
    }

    // warehouse is invalid
    props = setCommonConnectionParameters(true);
    props.put("warehouse", "NOT_EXISTS");
    try
    {
      DriverManager.getConnection(params.get("uri"), props);
      fail("should fail");
    }
    catch (SQLException ex)
    {
      assertEquals("error code", ex.getErrorCode(), SESSION_CREATION_OBJECT_DOES_NOT_EXIST_NOT_AUTHORIZED);
    }

    // role is invalid
    props = setCommonConnectionParameters(true);
    props.put("role", "NOT_EXISTS");
    try
    {
      DriverManager.getConnection(params.get("uri"), props);
      fail("should fail");
    }
    catch (SQLException ex)
    {
      assertEquals("error code", ex.getErrorCode(), ROLE_IN_CONNECT_STRING_DOES_NOT_EXIST);
    }
  }

  @Test
  public void testNoValidateDefaultParameters() throws Throwable
  {
    Map<String, String> params = getConnectionParameters();
    Properties props;

    props = setCommonConnectionParameters(false);
    props.put("db", "NOT_EXISTS");
    DriverManager.getConnection(params.get("uri"), props);

    // schema is invalid
    props = setCommonConnectionParameters(false);
    props.put("schema", "NOT_EXISTS");
    DriverManager.getConnection(params.get("uri"), props);

    // warehouse is invalid
    props = setCommonConnectionParameters(false);
    props.put("warehouse", "NOT_EXISTS");
    DriverManager.getConnection(params.get("uri"), props);

    // role is invalid
    props = setCommonConnectionParameters(false);
    props.put("role", "NOT_EXISTS");
    try
    {
      DriverManager.getConnection(params.get("uri"), props);
      fail("should fail");
    }
    catch (SQLException ex)
    {
      assertEquals("error code", ex.getErrorCode(), ROLE_IN_CONNECT_STRING_DOES_NOT_EXIST);
    }
  }

  private Properties setCommonConnectionParameters(boolean validateDefaultParameters)
  {
    Map<String, String> params = getConnectionParameters();
    Properties props = new Properties();
    props.put("validateDefaultParameters", validateDefaultParameters);
    props.put("account", params.get("account"));
    props.put("ssl", params.get("ssl"));
    props.put("role", params.get("role"));
    props.put("user", params.get("user"));
    props.put("password", params.get("password"));
    props.put("db", params.get("database"));
    props.put("schema", params.get("schema"));
    props.put("warehouse", params.get("warehouse"));
    return props;
  }


  private class ConcurrentConnections implements Runnable
  {
    Connection con = null;

    ConcurrentConnections()
    {
    }

    @Override
    public void run()
    {
      try
      {
        con = getConnection();
        con.createStatement().executeQuery(
            "select * from bigTable");
        con.close();

      }
      catch (SQLException ex)
      {
        try
        {
          con.close();
        }
        catch (SQLException e)
        {
          e.printStackTrace();
        }
        ex.printStackTrace();
      }
    }
  }

  @Test
  public void testAsyncQueryOpenAndCloseConnection() throws SQLException, IOException, InterruptedException
  {
    // open connection and run asynchronous query
    Connection con = getConnection();
    Statement statement = con.createStatement();
    ResultSet rs1 =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 40))");
    // Retrieve query ID for part 2 of test, check status of query
    String queryID = rs1.unwrap(SnowflakeResultSet.class).getQueryID();
    Thread.sleep(100);
    QueryStatus status = rs1.unwrap(SnowflakeResultSet.class).getStatus();
    // Query should take 60 seconds so should be running
    assertEquals(RUNNING, status);
    // close connection and wait for 1 minute while query finishes running
    statement.close();
    con.close();
    Thread.sleep(1000 * 70);
    //Create a new connection and new instance of a resultSet using query ID
    con = getConnection();
    try
    {
      ResultSet rs = con.unwrap(SnowflakeConnection.class).createResultSet("Totally invalid query ID");
      fail("Query ID should be rejected");
    }
    catch (SQLException e)
    {
      assertEquals(SqlState.INVALID_PARAMETER_VALUE, e.getSQLState());
    }
    ResultSet rs = con.unwrap(SnowflakeConnection.class).createResultSet(queryID);
    status = rs.unwrap(SnowflakeResultSet.class).getStatus();
    // Assert status of query is a success
    assertEquals(QueryStatus.SUCCESS, status);
    assertEquals("No error reported", status.getErrorMessage());
    assertEquals(0, status.getErrorCode());
    assertEquals(1, getSizeOfResultSet(rs));
    statement = con.createStatement();
    // Create another query that will not be successful (querying table that does not exist)
    rs1 = statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select * from nonexistentTable");
    Thread.sleep(100);
    status = rs1.unwrap(SnowflakeResultSet.class).getStatus();
    assertEquals(QueryStatus.FAILED_WITH_ERROR, status);
    assertEquals(2003, status.getErrorCode());
    assertEquals("SQL compilation error:\n" +
                 "Object 'NONEXISTENTTABLE' does not exist or not " +
                 "authorized.", status.getErrorMessage());
    statement.close();
    con.close();
  }

  @Test
  public void testAsyncAndSynchronousQueries() throws SQLException
  {
    Connection con = getConnection();
    Statement statement = con.createStatement();
    // execute some statements that you want to be synchronous
    statement.execute("alter session set CLIENT_TIMESTAMP_TYPE_MAPPING=TIMESTAMP_TZ");
    statement.execute("create or replace table smallTable (colA string, colB int)");
    statement.execute("create or replace table uselessTable (colA string, colB int)");
    statement.execute("insert into smallTable values ('row1', 1), ('row2', 2), ('row3', 3)");
    statement.execute("insert into uselessTable values ('row1', 1), ('row2', 2), ('row3', 3)");
    // Select from uselessTable asynchronously; drop it synchronously afterwards
    ResultSet rs = statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select * from smallTable");
    // execute a query that you don't want to wait for
    ResultSet rs1 = statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select * from uselessTable");
    // Drop the table that was queried asynchronously. Should not drop until after async query finishes, because this
    // query IS synchronous
    ResultSet rs2 =
        statement.executeQuery("drop table uselessTable");
    while (rs2.next())
    {
      assertEquals("USELESSTABLE successfully dropped.", rs2.getString(1));
    }
    // able to successfully fetch results in spite of table being dropped
    assertEquals(3, getSizeOfResultSet(rs1));
    statement.execute("alter session set CLIENT_TIMESTAMP_TYPE_MAPPING=TIMESTAMP_LTZ");

    // come back to the asynchronously executed result set after finishing other things
    rs.next();
    assertEquals(rs.getString(1), "row1");
    assertEquals(rs.getInt(2), 1);
    rs.next();
    assertEquals(rs.getString(1), "row2");
    assertEquals(rs.getInt(2), 2);
    rs.next();
    assertEquals(rs.getString(1), "row3");
    assertEquals(rs.getInt(2), 3);
    statement.execute("drop table smallTable");
    statement.close();
    con.close();
  }

  @Test
  public void testPreparedStatementAsyncQuery() throws SQLException
  {
    Connection con = getConnection();
    con.createStatement().execute("create or replace table testTable(colA string, colB boolean)");
    PreparedStatement prepStatement = con.prepareStatement("insert into testTable values (?,?)");
    prepStatement.setInt(1, 33);
    prepStatement.setBoolean(2, true);
    // call executeAsyncQuery
    ResultSet rs = prepStatement.unwrap(SnowflakePreparedStatement.class).executeAsyncQuery();
    // Get access to results by calling next() function
    // next () will block until results are ready
    assertTrue(rs.next());
    // the resultSet consists of a single row, single column containing the number of rows that have been updated by the insert
    // the number of updated rows in testTable is 1 so 1 is returned
    assertEquals(rs.getString(1), "1");
    con.createStatement().execute("drop table testTable");
    prepStatement.close();
    con.close();
  }

  @Test
  public void testIsStillRunning()
  {
    QueryStatus[] runningStatuses = {QueryStatus.RUNNING, QueryStatus.RESUMING_WAREHOUSE, QueryStatus.QUEUED,
                                     QueryStatus.QUEUED_REPAIRING_WAREHOUSE, QueryStatus.NO_DATA};

    QueryStatus[] otherStatuses = {QueryStatus.ABORTED, QueryStatus.ABORTING, QueryStatus.SUCCESS,
                                   QueryStatus.FAILED_WITH_ERROR, QueryStatus.FAILED_WITH_INCIDENT,
                                   QueryStatus.DISCONNECTED, QueryStatus.RESTARTED, QueryStatus.BLOCKED};

    for (QueryStatus qs : runningStatuses)
    {
      assertEquals(true, QueryStatus.isStillRunning(qs));
    }

    for (QueryStatus qs : otherStatuses)
    {
      assertEquals(false, QueryStatus.isStillRunning(qs));
    }
  }

  @Test
  public void testIsAnError()
  {
    QueryStatus[] otherStatuses = {QueryStatus.RUNNING, QueryStatus.RESUMING_WAREHOUSE, QueryStatus.QUEUED,
                                   QueryStatus.QUEUED_REPAIRING_WAREHOUSE, QueryStatus.SUCCESS, QueryStatus.RESTARTED,
                                   QueryStatus.NO_DATA};

    QueryStatus[] errorStatuses = {QueryStatus.ABORTED, QueryStatus.ABORTING,
                                   QueryStatus.FAILED_WITH_ERROR, QueryStatus.FAILED_WITH_INCIDENT,
                                   QueryStatus.DISCONNECTED, QueryStatus.BLOCKED};

    for (QueryStatus qs : errorStatuses)
    {
      assertEquals(true, QueryStatus.isAnError(qs));
    }

    for (QueryStatus qs : otherStatuses)
    {
      assertEquals(false, QueryStatus.isAnError(qs));
    }
  }


  /**
   * MANUAL TESTING OF ASYNCHRONOUS QUERYING
   * <p>
   * This test does not provide reliable results because the status of the queries
   * often depends on the GS server's behavior. We can often replicate
   * QUEUED and RESUMING_WAREHOUSE statuses, however.
   */
  //@Test
  public void testQueryStatuses()
  throws SQLException, IOException, InterruptedException
  {
    // Before running test, close warehouse and re-open it!
    Connection con = getConnection();
    Statement statement = con.createStatement();
    ResultSet rs =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 5))");
    Thread.sleep(100);
    QueryStatus status = rs.unwrap(SnowflakeResultSet.class).getStatus();
    // Since warehouse has just been restarted, warehouse should still be booting
    assertEquals(QueryStatus.RESUMING_WAREHOUSE, status);

    // now try to get QUEUED status
    ResultSet rs1 =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 60))");
    ResultSet rs2 =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 60))");
    ResultSet rs3 =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 60))");
    ResultSet rs4 =
        statement.unwrap(SnowflakeStatement.class).executeAsyncQuery("select count(*) from table(generator(timeLimit => 60))");
    // Retrieve query ID for part 2 of test, check status of query
    Thread.sleep(100);
    status = rs4.unwrap(SnowflakeResultSet.class).getStatus();
    // Since 4 queries were started at once, status is most likely QUEUED
    assertEquals(QueryStatus.QUEUED, status);
  }
}