net.snowflake.client.ingest.IngestTester Maven / Gradle / Ivy

Go to download
/*
 * Copyright (c) 2012-2016 Snowflake Computing Inc. All right reserved.
 */

package net.snowflake.client.ingest;

import net.snowflake.client.IngestFilesTester;
import net.snowflake.client.TestConnectionUtil;
import org.apache.commons.io.FileUtils;
import org.apache.http.impl.client.CloseableHttpClient;

import java.io.File;
import java.io.IOException;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ThreadLocalRandom;

/**
 * Created by hyu on 10/19/16.
 */
public class IngestTester extends IngestFilesTester
{
  final String user = "user" + RND_PAD;
  final String truthTable = INGEST_TBL + "_TRUTH";
  public final String insert_role = "API_INSERT";
  public final String CREATE_STAGE = "create or replace stage " + quote(STAGE_NAME)
                                   + " url='file://" + tmpIngestStage + "'";

  public final String CREATE_STAGE_S3 = "create or replace stage " + quote(STAGE_NAME)
                                    +" url='s3://sfc-dev1-data/hyu/ingest/'";
  public final String CREATE_STAGE_S3_MK = "create or replace stage " + quote
    (STAGE_NAME) +" url='s3://sfc-dev1-data/tpch/sf1_encrypted/region/'";

  public final String CREATE_INT_STAGE_S3 = "create or replace stage " +
    quote(STAGE_NAME);

  public final String CREATE_INGEST_TABLE = "create or replace table " + quote(INGEST_TBL)
    + " (row_id int, row_str string, num int, src string)";

  public final String CREATE_INGEST_MK_TABLE = "create or replace table " +
    quote(INGEST_TBL) + "(r_regionkey bigint, r_name char(25), r_comment varchar(152))";

  public final String CREATE_INGEST_TABLE_TRUTH = "create or replace table " + quote(truthTable)
                                + " (row_id int, row_str string, num int, src string)";

  public final String CREATE_INGEST_TABLE_VAR = "create or replace table " + quote(INGEST_TBL)
                  + "(v variant)";

  public final String CREATE_INGEST_TABLE_VAR_TRUTH = "create or replace table "  + quote(truthTable)
                  + "(v variant)";

  private String jwtToken;

  public IngestTester(Connection connection, CloseableHttpClient httpClient)
  {
    super(connection, httpClient);
  }

  public void testSetUp() throws Exception
  {
    doQuery("create database " + DATABASE);
    doQuery("use database " + DATABASE);
    doQuery("create schema " + SCHEMA);
    doQuery("use schema " + SCHEMA);

    doQuery("use role accountadmin");
    doQuery("create or replace role " + insert_role);

    jwtToken = createUserAndAuthenticate(user);
  }

  public String createUserAndAuthenticate(String userName) throws Exception
  {
    doQuery("create user " + userName + " default_role = " + insert_role +
                " password=\"" + TestConnectionUtil.PASSWORD + "\"");
    doQuery("grant role " + insert_role + " to user " + userName);

    final String account = TestConnectionUtil.ACCOUNT;

    say("begin setup security");
    final SecurityState securityState = SecurityState.create(account, userName);
    say("end setup security");

    String pk = securityState.getPublicKeyString();
    assert pk != null;
    assert !pk.isEmpty();

    doQuery("alter user " + userName +
        " set RSA_PUBLIC_KEY='" + pk + "'");

    return securityState.getJwtToken();
  }

  public void grantPriviOnTableToRole() throws SQLException
  {
    doQuery("grant INSERT on " + quote(INGEST_TBL) + " to ROLE " + insert_role);

    //for read history.
    doQuery("grant SELECT on " + quote(INGEST_TBL) + " to ROLE " + insert_role);

    //TODO: avoid grant ALL - yet we need to be able to create the pipe...
    doQuery("grant ALL on DATABASE " + DATABASE +
        " to ROLE " + insert_role);
    doQuery("grant ALL on SCHEMA " + SCHEMA +
        " to ROLE " + insert_role);

    doQuery("grant USAGE on STAGE " + quote(STAGE_NAME)
        + " to ROLE " + insert_role);

  }

  public void grantPriviOnTableToRoleInternal() throws SQLException
  {
    doQuery("grant INSERT on " + quote(INGEST_TBL) + " to ROLE " + insert_role);

    //for read history.
    doQuery("grant SELECT on " + quote(INGEST_TBL) + " to ROLE " + insert_role);

    //TODO: avoid grant ALL - yet we need to be able to create the pipe...
    doQuery("grant ALL on DATABASE " + DATABASE +
      " to ROLE " + insert_role);
    doQuery("grant ALL on SCHEMA " + SCHEMA +
      " to ROLE " + insert_role);

    doQuery("grant READ on STAGE " + quote(STAGE_NAME)
      + " to ROLE " + insert_role);

    doQuery("grant WRITE on STAGE " + quote(STAGE_NAME)
      + " to ROLE " + insert_role);

  }

  public void createTableWithWideColumns(int columnNum) throws Exception
  {
    StringBuilder createIngestTableStr = new StringBuilder(
        "create or replace table " + quote(INGEST_TBL) + "(");
    StringBuilder createIngestTableTruthStr = new StringBuilder(
        "create or replace table " + quote(truthTable) + "(");
    for (int i=0; i createMalformedCSVs(int numOfFiles, int rows) throws IOException
  {
    ThreadLocalRandom rnd = ThreadLocalRandom.current();
    int malformedFileIndex = rnd.nextInt(numOfFiles);
    List files = new ArrayList<>();
    for (int i=0; i< numOfFiles; i++)
    {
      if (i == malformedFileIndex)
        files.add(createMalformedCSV(rows));
      else
        files.add(createTempCsv(rows));
    }
    return files;
  }

  public List createCsvs(int numOfFiles, int rows) throws IOException
  {
    List files = new ArrayList<>();
    for (int i = 0; i < numOfFiles; i++)
    {
      files.add(createTempCsv(rows));
    }
    return files;
  }

  public Path createLargeCSV(int rows, int numCol)
      throws IOException
  {
    final Path csv = Files.createTempFile(tmpIngestStage, "data", ".csv");
    try (Writer w = Files.newBufferedWriter(csv, StandardCharsets.UTF_8))
    {
      for (int i = 0; i < rows; i++)
      {
        for (int j=0; j