Please wait. This can take some minutes ...
Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
net.snowflake.client.loader.test.LoaderIT Maven / Gradle / Ivy
/*
* Copyright (c) 2012-2016 Snowflake Computing Inc. All right reserved.
*/
package net.snowflake.client.loader.test;
import net.snowflake.client.jdbc.SnowflakeConnectionV1;
import net.snowflake.client.loader.LoadResultListener;
import net.snowflake.client.loader.Loader;
import net.snowflake.client.loader.LoaderFactory;
import net.snowflake.client.loader.LoaderProperty;
import net.snowflake.client.loader.LoadingError;
import net.snowflake.client.loader.StreamLoader;
import net.snowflake.client.loader.Operation;
import net.snowflake.client.loader.BufferStage;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.sql.PreparedStatement;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.TimeZone;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.CoreMatchers.*;
import static org.junit.Assert.*;
/**
* Loader IT
*/
public class LoaderIT
{
public final static String TARGET_TABLE_NAME = "LOADER_test_TABLE";
static StreamLoader underTest;
static Connection testConnection;
static Connection putConnection;
public LoaderIT()
{
}
@BeforeClass
public static void setUpClass() throws Throwable
{
Class.forName("net.snowflake.client.jdbc.SnowflakeDriver");
String accountStr = "testaccount";
if (System.getenv("SF_ACCOUNT") != null)
{
accountStr = System.getenv("SF_ACCOUNT");
}
// create a new connection
String connectionUrl = System.getenv("SF_JDBC_CONNECT_STRING");
// use the default connection string if it is not set in environment
if (connectionUrl == null)
{
connectionUrl = "jdbc:snowflake://localhost:8080";
}
connectionUrl = connectionUrl
+ "/?account=" + accountStr
+ "&user=snowman&password=test&schema=testschema&db=testdb";
Properties properties = new Properties();
properties.put("internal", "true");
properties.put("ssl", "off");
testConnection = DriverManager.getConnection(connectionUrl, properties);
putConnection = DriverManager.getConnection(connectionUrl, properties);
testConnection.createStatement().execute(
"CREATE OR REPLACE STAGE loadertest URL='file:///tmp/loadertest'");
testConnection.createStatement().execute(
"CREATE OR REPLACE SCHEMA LOADER");
testConnection.createStatement().execute(
"USE SCHEMA testschema");
testConnection.createStatement().execute(String.format(
"CREATE OR REPLACE TABLE LOADER.\"%s\" ("
+ "ID int, "
+ "C1 varchar(255), "
+ "C2 varchar(255) DEFAULT 'X', "
+ "C3 double, "
+ "C4 timestamp, "
+ "C5 variant)", LoaderIT.TARGET_TABLE_NAME));
}
@AfterClass
public static void tearDownClass() throws SQLException
{
testConnection.createStatement().execute(
String.format("DROP TABLE IF EXISTS LOADER.\"%s\"", LoaderIT.TARGET_TABLE_NAME));
testConnection.createStatement().execute(
"DROP SCHEMA IF EXISTS LOADER");
testConnection.createStatement().execute(
"DROP STAGE IF EXISTS loadertest");
testConnection.close();
putConnection.close();
}
@Before
public void setUp() {
}
@After
public void tearDown() {
}
private Map initLoaderProperties() throws Exception {
return initLoaderProperties(TARGET_TABLE_NAME);
}
private Map initLoaderProperties(String tableName) throws Exception
{
Map prop = new HashMap<>();
prop.put(LoaderProperty.tableName, tableName);
prop.put(LoaderProperty.schemaName, "LOADER");
prop.put(LoaderProperty.databaseName, "TESTDB");
prop.put(LoaderProperty.columns, Arrays.asList(
"ID"
));
prop.put(LoaderProperty.remoteStage, "loadertest");
prop.put(LoaderProperty.operation, Operation.INSERT);
prop.put(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C3", "C4", "C5"
));
return prop;
}
private ResultListener initLoader(
Map prop,
boolean testMode,
long csvFileBucketSize, long csvFileSize) throws Exception
{
ResultListener _resultListener = new ResultListener();
// Delete staging area
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.startTransaction, true);
underTest.setProperty(LoaderProperty.truncateTable, true);
underTest.setProperty(LoaderProperty.executeBefore,
"CREATE TABLE DUMMY_TABLE(i int)");
underTest.setProperty(LoaderProperty.executeAfter,
"DROP TABLE DUMMY_TABLE");
underTest.start();
underTest.finish();
// Set up Test parameters
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
// file bucket size
underTest.setProperty(LoaderProperty.csvFileBucketSize,
Long.toString(csvFileBucketSize));
// file batch
underTest.setProperty(LoaderProperty.csvFileSize,
Long.toString(csvFileSize));
underTest.setListener(_resultListener);
// causes upload to fail
underTest.setTestMode(testMode);
// Wait for 5 seconds on first put to buffer everything up.
((SnowflakeConnectionV1) putConnection).setInjectedDelay(5000);
return _resultListener;
}
private List populateTestData(
boolean testMode) throws Exception
{
return populateTestData(testMode,
null,
TARGET_TABLE_NAME,
10000, BufferStage.FILE_BUCKET_SIZE, BufferStage.FILE_SIZE);
}
private List populateTestData(
boolean testMode,
List dataSet,
String tableName,
int numberOfRows, long csvFileBucketSize, long csvFileSize) throws Exception
{
Map prop = this.initLoaderProperties(tableName);
ResultListener listener = this.initLoader(prop, testMode, csvFileBucketSize, csvFileSize);
underTest.start();
Random rnd = new Random();
List newDataSet = new ArrayList<>();
if (dataSet == null) {
// generates a new data set and ingest
for (int i = 0; i < numberOfRows; i++) {
final String json = "{\"key\":" + String.valueOf(rnd.nextInt()) + ","
+ "\"bar\":" + i + "}";
Object[] row = new Object[]
{
i, "foo_" + i, rnd.nextInt() / 3, new Date(),
json
};
newDataSet.add(row);
underTest.submitRow(row);
}
} else {
for (Object[] row: dataSet) {
// ingest the same data set
underTest.submitRow(row);
}
}
underTest.finish();
int submitted = listener.getSubmittedRowCount();
assertThat("submitted rows",
submitted, equalTo(numberOfRows));
assertThat("_resultListener.counter is not correct",
listener.counter.get(), equalTo(numberOfRows));
assertThat("_resultListener.getErrors() was not 0",
listener.getErrors().size(), equalTo(0));
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT COUNT(*) AS N"
+ " FROM LOADER.\"%s\"", tableName));
rs.next();
int count = rs.getInt("N");
assertThat("count is not correct", count, equalTo(numberOfRows));
assertThat("_resultListener.processed didn't match count",
listener.processed.get(), equalTo(count));
assertThat("_resultListener.counter didn't match count",
listener.counter.get(), equalTo(count));
assertThat("_resultListener.getErrors().size() was not 0",
listener.getErrors().size(), equalTo(0));
assertThat("_resultListener.getLastRecord()[0] was not 9999",
(Integer) listener.getLastRecord()[0], equalTo(numberOfRows-1));
return newDataSet;
}
@Test
public void testInjectBadStagedFileInsert() throws Exception {
Map prop = this.initLoaderProperties(TARGET_TABLE_NAME);
ResultListener listener = this.initLoader(
prop, false, BufferStage.FILE_BUCKET_SIZE, BufferStage.FILE_SIZE);
listener.throwOnError = true;
int numberOfRows = 1000;
underTest.setProperty(LoaderProperty.testRemoteBadCSV, true);
underTest.setProperty(LoaderProperty.startTransaction, true);
underTest.start();
Random rnd = new Random();
List newDataSet = new ArrayList<>();
// generates a new data set and ingest
for (int i = 0; i < numberOfRows; i++) {
final String json = "{\"key\":" + String.valueOf(rnd.nextInt()) + ","
+ "\"bar\":" + i + "}";
Object[] row = new Object[]
{
i, "foo_" + i, rnd.nextInt() / 3, new Date(),
json
};
newDataSet.add(row);
underTest.submitRow(row);
}
try {
underTest.finish();
fail("Should raise and error");
} catch(Loader.DataError ex) {
assertThat("Loader.DataError is raised", true);
}
}
@Test
public void testExecuteBeforeAfterSQLError() throws Exception {
Map prop = this.initLoaderProperties(
TARGET_TABLE_NAME);
this.initLoader(
prop, false, BufferStage.FILE_BUCKET_SIZE, BufferStage.FILE_SIZE);
underTest.setProperty(
LoaderProperty.executeBefore, "SELECT * FROOOOOM TBL");
underTest.start();
try {
underTest.finish();
fail("SQL Error should be raised.");
} catch (Loader.ConnectionError e) {
assertThat(e.getCause(), instanceOf(SQLException.class));
}
this.initLoader(
prop, false, BufferStage.FILE_BUCKET_SIZE, BufferStage.FILE_SIZE);
underTest.setProperty(
LoaderProperty.executeBefore, "select current_version()");
underTest.setProperty(
LoaderProperty.executeAfter, "SELECT * FROM TBBBBBBL");
underTest.start();
try {
underTest.finish();
fail("SQL Error should be raised.");
} catch (Loader.ConnectionError e) {
assertThat(e.getCause(), instanceOf(SQLException.class));
}
}
@Test
public void testLoaderInsert() throws Exception
{
// mostly just populate test data but with delay injection to test
// PUT retry
this.populateTestData(true);
}
@Test
public void testLoaderMultipleBatch() throws Exception {
String refTableName = "LOADER_TEST_TABLE_REF";
testConnection.createStatement().execute(String.format(
"CREATE OR REPLACE TABLE LOADER.%s ("
+ "ID int, "
+ "C1 varchar(255), "
+ "C2 varchar(255) DEFAULT 'X', "
+ "C3 double, "
+ "C4 timestamp, "
+ "C5 variant)", refTableName));
try {
List dataSet = this.populateTestData(
false,
null,
TARGET_TABLE_NAME,
10000, 64, 5000000);
this.populateTestData(
false,
dataSet,
refTableName,
10000, 2, 30000);
ResultSet rsReference = testConnection.createStatement().executeQuery(String.format(
"SELECT hash_agg(*) FROM LOADER.\"%s\"", TARGET_TABLE_NAME
));
rsReference.next();
long hashValueReference = rsReference.getLong(1);
ResultSet rsTarget = testConnection.createStatement().executeQuery(String.format(
"SELECT hash_agg(*) FROM LOADER.\"%s\"", refTableName
));
rsTarget.next();
long hashValueTarget = rsTarget.getLong(1);
assertThat("hash values", hashValueTarget, equalTo(hashValueReference));
} finally {
testConnection.createStatement().execute(String.format(
"DROP TABLE IF EXISTS %s", refTableName));
}
}
@Test
public void testLoaderDelete() throws Exception
{
this.populateTestData(false);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID", "C1"
));
underTest.setProperty(LoaderProperty.operation, Operation.DELETE);
underTest.setListener(new ResultListener());
underTest.start();
Object[] del = new Object[]
{
42, "foo_42" // deleted
};
underTest.submitRow(del);
del = new Object[]
{
41, "blah" // ignored, and should not raise any error/warning
};
underTest.submitRow(del);
underTest.finish();
ResultListener listener = (ResultListener) underTest.getListener();
assertThat("error count", listener.getErrorCount(), equalTo(0));
assertThat("error record count",
listener.getErrorRecordCount(), equalTo(0));
assertThat("submitted row count",
listener.getSubmittedRowCount(), equalTo(2));
assertThat("processed", listener.processed.get(), equalTo(1));
assertThat("deleted rows", listener.deleted.get(), equalTo(1));
}
@Test
public void testLoaderModify() throws Exception
{
this.populateTestData(false); // this does INSERT operation
Map prop = this.initLoaderProperties();
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2", "C3", "C4", "C5"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID"
));
underTest.setProperty(LoaderProperty.operation, Operation.MODIFY);
underTest.setListener(new ResultListener());
underTest.start();
Object[] mod = new Object[]
{
41, "modified", "some\nthi\"ng\\", 41.6, new Date(), "{}"
};
underTest.submitRow(mod);
mod = new Object[]
{
40, "modified", "\"something,", 40.2, new Date(), "{}"
};
underTest.submitRow(mod);
underTest.finish();
ResultListener listener = (ResultListener) underTest.getListener();
assertThat("processed", listener.processed.get(), equalTo(2));
assertThat("submitted row", listener.getSubmittedRowCount(), equalTo(2));
assertThat("updated", listener.updated.get(), equalTo(2));
assertThat("error count", listener.getErrorCount(), equalTo(0));
assertThat("error record count", listener.getErrorRecordCount(), equalTo(0));
// Test deletion
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT COUNT(*) AS N"
+ " FROM LOADER.\"%s\"", TARGET_TABLE_NAME));
rs.next();
assertThat("count is not correct", rs.getInt("N"), equalTo(10000));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1 AS N"
+ " FROM LOADER.\"%s\" WHERE ID=40", TARGET_TABLE_NAME));
rs.next();
assertThat("status is not correct", rs.getString("N"), equalTo("modified"));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1, C2"
+ " FROM LOADER.\"%s\" WHERE ID=41", TARGET_TABLE_NAME));
rs.next();
assertThat("C1 is not correct",
rs.getString("C1"), equalTo("modified"));
assertThat("C2 is not correct",
rs.getString("C2"), equalTo("some\nthi\"ng\\"));
}
@Test
public void testLoaderModifyWithOneMatchOneNot() throws Exception
{
this.populateTestData(false); // this does INSERT operation
Map prop = this.initLoaderProperties();
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2", "C3", "C4", "C5"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID"
));
underTest.setProperty(LoaderProperty.operation, Operation.MODIFY);
underTest.setListener(new ResultListener());
underTest.start();
Object[] mod = new Object[]
{
20000, "modified", "some\nthi\"ng\\", 41.6, new Date(), "{}"
};
underTest.submitRow(mod);
mod = new Object[]
{
45, "modified", "\"something2,", 40.2, new Date(), "{}"
};
underTest.submitRow(mod);
underTest.finish();
ResultListener listener = (ResultListener) underTest.getListener();
assertThat("processed", listener.processed.get(), equalTo(1));
assertThat("submitted row", listener.getSubmittedRowCount(), equalTo(2));
assertThat("updated", listener.updated.get(), equalTo(1));
assertThat("error count", listener.getErrorCount(), equalTo(0));
assertThat("error record count", listener.getErrorRecordCount(), equalTo(0));
// Test deletion
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT COUNT(*) AS N"
+ " FROM LOADER.\"%s\"", TARGET_TABLE_NAME));
rs.next();
assertThat("count is not correct", rs.getInt("N"), equalTo(10000));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1, C2"
+ " FROM LOADER.\"%s\" WHERE ID=45", TARGET_TABLE_NAME));
rs.next();
assertThat("C1 is not correct",
rs.getString("C1"), equalTo("modified"));
assertThat("C2 is not correct",
rs.getString("C2"), equalTo("\"something2,"));
}
@Test
public void testLoaderUpsert() throws Exception
{
this.populateTestData(false);
Map prop = this.initLoaderProperties();
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2", "C3", "C4", "C5"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID"
));
underTest.setProperty(LoaderProperty.operation, Operation.UPSERT);
underTest.setListener(new ResultListener());
underTest.start();
Date d = new Date();
Object[] ups = new Object[]
{
10001, "inserted\\,", "something", 0x4.11_33p2, d, "{}"
};
underTest.submitRow(ups);
ups = new Object[]
{
39, "modified", "something", 40.1, d, "{}"
};
underTest.submitRow(ups);
underTest.finish();
ResultListener listener = (ResultListener) underTest.getListener();
assertThat("processed", listener.processed.get(), equalTo(2));
assertThat("submitted row", listener.getSubmittedRowCount(), equalTo(2));
assertThat("updated/inserted", listener.updated.get(), equalTo(2));
assertThat("error count", listener.getErrorCount(), equalTo(0));
assertThat("error record count", listener.getErrorRecordCount(), equalTo(0));
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1, C4, C3"
+ " FROM LOADER.\"%s\" WHERE ID=10001", TARGET_TABLE_NAME));
rs.next();
assertThat("C1 is not correct", rs.getString("C1"), equalTo("inserted\\,"));
long l = rs.getTimestamp("C4").getTime();
assertThat("C4 is not correct", l, equalTo(d.getTime()));
assertThat("C3 is not correct", Double.toHexString((rs.getDouble("C3"))),
equalTo("0x1.044cc0000225cp4"));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1 AS N"
+ " FROM LOADER.\"%s\" WHERE ID=39", TARGET_TABLE_NAME));
rs.next();
assertThat("N is not correct", rs.getString("N"), equalTo("modified"));
}
@Test
public void testLoaderUpsertWithError() throws Exception
{
this.populateTestData(false);
Map prop = this.initLoaderProperties();
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2", "C3", "C4", "C5"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID"
));
underTest.setProperty(LoaderProperty.operation, Operation.UPSERT);
underTest.setListener(new ResultListener());
underTest.start();
underTest.setListener(new ResultListener());
Object[] upse = new Object[]
{
"10001-", "inserted", "something", "42-", new Date(), "{}"
};
underTest.submitRow(upse);
upse = new Object[]
{
10002, "inserted", "something", 43, new Date(), "{}"
};
underTest.submitRow(upse);
upse = new Object[]
{
45, "modified", "something", 46.1, new Date(), "{}"
};
underTest.submitRow(upse);
underTest.finish();
ResultListener listener = (ResultListener) underTest.getListener();
assertThat("processed", listener.processed.get(), equalTo(3));
assertThat("counter", listener.counter.get(), equalTo(2));
assertThat("submitted row", listener.getSubmittedRowCount(), equalTo(3));
assertThat("updated/inserted", listener.updated.get(), equalTo(2));
assertThat("error count", listener.getErrorCount(), equalTo(2));
assertThat("error record count", listener.getErrorRecordCount(), equalTo(1));
assertThat("Target table name is not correct",listener.getErrors().get(0)
.getTarget(), equalTo(TARGET_TABLE_NAME));
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT COUNT(*) AS N"
+ " FROM LOADER.\"%s\"", TARGET_TABLE_NAME));
rs.next();
int c = rs.getInt("N");
assertThat("N is not correct", c, equalTo(10001));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C1 AS N"
+ " FROM LOADER.\"%s\" WHERE ID=45", TARGET_TABLE_NAME));
rs.next();
assertThat("N is not correct", rs.getString("N"), equalTo("modified"));
}
@Test
public void testLoaderUpsertWithErrorAndRollback() throws Exception
{
this.populateTestData(false);
PreparedStatement pstmt = testConnection.prepareStatement(
String.format("INSERT INTO LOADER.\"%s\"(ID,C1,C2,C3,C4,C5)"
+ " SELECT column1, column2, column3, column4,"
+ " column5, parse_json(column6)"
+ " FROM VALUES(?,?,?,?,?,?)", TARGET_TABLE_NAME));
pstmt.setInt(1, 10001);
pstmt.setString(2, "inserted\\,");
pstmt.setString(3, "something");
pstmt.setDouble(4, 0x4.11_33p2);
pstmt.setDate(5, new java.sql.Date(new Date().getTime()));
pstmt.setObject(6, "{}");
pstmt.execute();
testConnection.commit();
Map prop = this.initLoaderProperties();
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2", "C3", "C4", "C5"
));
underTest.setProperty(LoaderProperty.keys, Arrays.asList(
"ID"
));
underTest.setProperty(LoaderProperty.operation, Operation.UPSERT);
underTest.setProperty(LoaderProperty.startTransaction, true);
underTest.setProperty(LoaderProperty.preserveStageFile, true);
underTest.start();
ResultListener listener = new ResultListener();
listener.throwOnError = true; // should trigger rollback
underTest.setListener(listener);
try
{
Object[] noerr = new Object[]
{
"10001", "inserted", "something", "42", new Date(), "{}"
};
underTest.submitRow(noerr);
Object[] err = new Object[]
{
"10002-", "inserted", "something", "42-", new Date(), "{}"
};
underTest.submitRow(err);
underTest.finish();
fail("Test must raise Loader.DataError exception");
}
catch (Loader.DataError e)
{
// we are good
assertThat("error message",
e.getMessage(), allOf(
containsString("10002-"),
containsString("not recognized")));
}
assertThat("processed", listener.processed.get(), equalTo(0));
assertThat("submitted row", listener.getSubmittedRowCount(), equalTo(2));
assertThat("updated/inserted", listener.updated.get(), equalTo(0));
assertThat("error count", listener.getErrorCount(), equalTo(2));
assertThat("error record count", listener.getErrorRecordCount(), equalTo(1));
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT COUNT(*) AS N FROM LOADER.\"%s\"", TARGET_TABLE_NAME));
rs.next();
assertThat("N", rs.getInt("N"), equalTo(10001));
rs = testConnection.createStatement().executeQuery(
String.format("SELECT C3 FROM LOADER.\"%s\" WHERE id=10001", TARGET_TABLE_NAME));
rs.next();
assertThat("C3. No commit should happen",
Double.toHexString((rs.getDouble("C3"))),
equalTo("0x1.044cc0000225cp4"));
}
@Test
public void testLoadTimestamp() throws Exception {
String targetTableName = "LOADER_TEST_TIMESTAMP";
// init properties
Map prop = new HashMap<>();
prop.put(LoaderProperty.tableName, targetTableName);
prop.put(LoaderProperty.schemaName, "LOADER");
prop.put(LoaderProperty.databaseName, "TESTDB");
prop.put(LoaderProperty.columns, Arrays.asList(
"ID"
));
prop.put(LoaderProperty.remoteStage, "loadertest");
prop.put(LoaderProperty.operation, Operation.INSERT);
prop.put(LoaderProperty.columns, Arrays.asList(
"ID", "C1", "C2"
));
// create table including TIMESTAMP_NTZ
testConnection.createStatement().execute(String.format(
"CREATE OR REPLACE TABLE LOADER.%s ("
+ "ID int, "
+ "C1 varchar(255), "
+ "C2 timestamp_ntz)", targetTableName));
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
underTest.setProperty(LoaderProperty.startTransaction, true);
underTest.setProperty(LoaderProperty.truncateTable, true);
// file bucket size
underTest.setProperty(LoaderProperty.csvFileBucketSize,
Long.toString(BufferStage.FILE_BUCKET_SIZE));
// file batch
underTest.setProperty(LoaderProperty.csvFileSize,
Long.toString(BufferStage.FILE_SIZE));
// This test case is testing this flag
underTest.setProperty(LoaderProperty.useLocalTimezone, true);
ResultListener resultListener = new ResultListener();
underTest.setListener(resultListener);
underTest.start();
final TimeZone originalTimeZone = TimeZone.getDefault();
final TimeZone targetTimeZone = TimeZone.getTimeZone("America/Los_Angeles");
TimeZone.setDefault(targetTimeZone);
// input timestamp associated with the target timezone, America/Los_Angeles
Date currentTs = new Date();
try {
List newDataSet = new ArrayList<>();
for (int i = 0; i < 5; ++i) {
Object[] row = new Object[]{
i, "foo_" + i, currentTs
};
newDataSet.add(row);
underTest.submitRow(row);
}
underTest.finish();
} finally {
testConnection.createStatement().execute(String.format(
"DROP TABLE IF EXISTS %s", targetTableName));
TimeZone.setDefault(originalTimeZone);
}
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT * FROM LOADER.\"%s\"", targetTableName));
rs.next();
Timestamp ts = rs.getTimestamp("C2");
// format the input TS with the target timezone
SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS");
sdf.setTimeZone(targetTimeZone);
String currenTsStr = sdf.format(currentTs);
// format the retrieved TS with the original timezone
sdf.setTimeZone(originalTimeZone);
String retrievedTsStr = sdf.format(new Date(ts.getTime()));
// They must be identical.
assertThat("Input and retrieved timestamp are different",
retrievedTsStr, equalTo(currenTsStr));
}
// DISABLED @Test
public void testSpacesInColumnTable() throws Exception {
String targetTableName = "Load Test Spaces In Columns";
// init properties
Map prop = new HashMap<>();
prop.put(LoaderProperty.databaseName, "TESTDB");
prop.put(LoaderProperty.schemaName, "LOADER");
prop.put(LoaderProperty.tableName, targetTableName);
prop.put(LoaderProperty.columns, Arrays.asList(
"ID"
));
prop.put(LoaderProperty.remoteStage, "loadertest");
prop.put(LoaderProperty.operation, Operation.INSERT);
prop.put(LoaderProperty.columns, Arrays.asList(
"ID", "Column 1"
));
// create table with spaces in column names
testConnection.createStatement().execute(String.format(
"CREATE OR REPLACE TABLE \"LOADER\".\"%s\" ("
+ "ID int, "
+ "\"Column 1\" varchar(255))", targetTableName));
underTest = (StreamLoader) LoaderFactory.createLoader(
prop, putConnection, testConnection);
ResultListener resultListener = new ResultListener();
underTest.setListener(resultListener);
underTest.start();
List newDataSet = new ArrayList<>();
for (int i = 0; i < 5; ++i) {
Object[] row = new Object[]{
i, "foo_" + i
};
newDataSet.add(row);
underTest.submitRow(row);
}
underTest.finish();
ResultSet rs = testConnection.createStatement().executeQuery(
String.format("SELECT * FROM LOADER.\"%s\" ORDER BY \"Column 1\"",
targetTableName));
rs.next();
assertThat("The first id", rs.getInt(1), equalTo(0));
assertThat("The first str", rs.getString(2), equalTo("foo_0"));
}
class ResultListener implements LoadResultListener
{
final private List errors = new ArrayList<>();
final private AtomicInteger errorCount = new AtomicInteger(0);
final private AtomicInteger errorRecordCount = new AtomicInteger(0);
final public AtomicInteger counter = new AtomicInteger(0);
final public AtomicInteger processed = new AtomicInteger(0);
final public AtomicInteger deleted = new AtomicInteger(0);
final public AtomicInteger updated = new AtomicInteger(0);
final private AtomicInteger submittedRowCount = new AtomicInteger(0);
private Object[] lastRecord = null;
public boolean throwOnError = false; // should not trigger rollback
@Override
public boolean needErrors()
{
return true;
}
@Override
public boolean needSuccessRecords()
{
return true;
}
@Override
public void addError(LoadingError error)
{
errors.add(error);
}
@Override
public boolean throwOnError()
{
return throwOnError;
}
public List getErrors()
{
return errors;
}
@Override
public void recordProvided(Operation op, Object[] record)
{
lastRecord = record;
}
@Override
public void addProcessedRecordCount(Operation op, int i)
{
processed.addAndGet(i);
}
@Override
public void addOperationRecordCount(Operation op, int i)
{
counter.addAndGet(i);
if (op == Operation.DELETE)
{
deleted.addAndGet(i);
}
else if (op == Operation.MODIFY || op == Operation.UPSERT)
{
updated.addAndGet(i);
}
}
public Object[] getLastRecord()
{
return lastRecord;
}
@Override
public int getErrorCount()
{
return errorCount.get();
}
@Override
public int getErrorRecordCount()
{
return errorRecordCount.get();
}
@Override
public void resetErrorCount()
{
errorCount.set(0);
}
@Override
public void resetErrorRecordCount()
{
errorRecordCount.set(0);
}
@Override
public void addErrorCount(int count)
{
errorCount.addAndGet(count);
}
@Override
public void addErrorRecordCount(int count)
{
errorRecordCount.addAndGet(count);
}
@Override
public void resetSubmittedRowCount()
{
submittedRowCount.set(0);
}
@Override
public void addSubmittedRowCount(int count)
{
submittedRowCount.addAndGet(count);
}
@Override
public int getSubmittedRowCount()
{
return submittedRowCount.get();
}
}
}