All Downloads are FREE. Search and download functionalities are using the official Maven repository.

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();
    }
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy