
net.snowflake.client.ingest.IngestTester Maven / Gradle / Ivy
/*
* 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
© 2015 - 2025 Weber Informatics LLC | Privacy Policy