• Home
• net.snowflake
• snowflake-jdbc
• 3.10.2
• source code
• SnowflakeFileTransferAgent.java

×

Project download

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.jdbc.SnowflakeFileTransferAgent Maven / Gradle / Ivy

/*
 * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
 */

package net.snowflake.client.jdbc;


import com.amazonaws.util.Base64;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Strings;
import com.google.common.io.ByteStreams;
import com.google.common.io.CountingOutputStream;
import net.snowflake.client.core.ObjectMapperFactory;
import net.snowflake.client.core.SFException;
import net.snowflake.client.core.SFFixedViewResultSet;
import net.snowflake.client.core.SFSession;
import net.snowflake.client.core.SFStatement;
import net.snowflake.client.jdbc.cloud.storage.SnowflakeStorageClient;
import net.snowflake.client.jdbc.cloud.storage.StageInfo;
import net.snowflake.client.jdbc.cloud.storage.StorageClientFactory;
import net.snowflake.client.jdbc.cloud.storage.StorageObjectMetadata;
import net.snowflake.client.jdbc.cloud.storage.StorageObjectSummary;
import net.snowflake.client.jdbc.cloud.storage.StorageObjectSummaryCollection;
import net.snowflake.client.jdbc.cloud.storage.StorageProviderException;
import net.snowflake.client.jdbc.telemetryOOB.TelemetryService;
import net.snowflake.client.log.ArgSupplier;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial;
import net.snowflake.common.core.SqlState;
import net.snowflake.common.util.ClassUtil;
import net.snowflake.common.util.FixedViewColumn;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.io.filefilter.WildcardFileFilter;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.DigestOutputStream;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.zip.GZIPOutputStream;

import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;

/**
 * Class for uploading/downloading files
 *
 * @author jhuang
 */
public class SnowflakeFileTransferAgent implements SnowflakeFixedView
{
  final static SFLogger logger =
      SFLoggerFactory.getLogger(SnowflakeFileTransferAgent.class);

  final static StorageClientFactory storageFactory = StorageClientFactory.getFactory();

  private static final ObjectMapper mapper =
      ObjectMapperFactory.getObjectMapper();

  // We will allow buffering of upto 128M data before spilling to disk during
  // compression and digest computation
  final static int MAX_BUFFER_SIZE = 1 << 27;
  public static final String SRC_FILE_NAME_FOR_STREAM = "stream";

  private static final String FILE_PROTOCOL = "file://";

  static private String localFSFileSep = systemGetProperty("file.separator");
  static private int DEFAULT_PARALLEL = 10;

  private String command;

  // list of files specified. Wildcard should be expanded already for uploading
  // For downloading, it the list of stage file names
  private Set sourceFiles;

  // big source files >=16MB, for which we will not upload them in serial mode
  // since TransferManager will parallelize upload
  private Set bigSourceFiles;

  // big source files < 16MB, for which we will upload them in paralle mode
  // with 4 threads by default
  private Set smallSourceFiles;

  static final private int BIG_FILE_THRESHOLD = 16 * 1024 * 1024;

  private Map fileMetadataMap;

  // stage related info
  private StageInfo stageInfo;

  // local location for where to download files to
  private String localLocation;

  private boolean showEncryptionParameter;

  // default parallelism
  private int parallel = DEFAULT_PARALLEL;

  private SFSession connection;
  private SFStatement statement;

  private InputStream sourceStream;
  private boolean sourceFromStream;
  private boolean compressSourceFromStream;

  private String destFileNameForStreamSource;

  public StageInfo getStageInfo()
  {
    return this.stageInfo;
  }

  // Encryption material
  private List encryptionMaterial;

  // Presigned URLs
  private List presignedUrls;

  // Index: Source file to encryption material
  HashMap srcFileToEncMat;

  // Index: Source file to presigned URL
  HashMap srcFileToPresignedUrl;

  public Map getStageCredentials()
  {
    return new HashMap<>(stageInfo.getCredentials());
  }

  public List getEncryptionMaterial()
  {
    return new ArrayList<>(encryptionMaterial);
  }

  public Map getSrcToMaterialsMap()
  {
    return new HashMap<>(srcFileToEncMat);
  }

  public Map getSrcToPresignedUrlMap()
  {
    return new HashMap<>(srcFileToPresignedUrl);
  }

  public String getStageLocation()
  {
    return stageInfo.getLocation();
  }

  private void initEncryptionMaterial(CommandType commandType, JsonNode jsonNode)
  throws SnowflakeSQLException, JsonProcessingException
  {
    encryptionMaterial = new ArrayList<>();
    JsonNode rootNode = jsonNode.path("data").path("encryptionMaterial");
    if (commandType == CommandType.UPLOAD)
    {
      logger.debug("initEncryptionMaterial: UPLOAD");

      RemoteStoreFileEncryptionMaterial encMat = null;
      if (!rootNode.isMissingNode() && !rootNode.isNull())
      {
        encMat = mapper.treeToValue(rootNode, RemoteStoreFileEncryptionMaterial.class);
      }
      encryptionMaterial.add(encMat);

    }
    else
    {
      logger.debug("initEncryptionMaterial: DOWNLOAD");

      if (!rootNode.isMissingNode() && !rootNode.isNull())
      {
        encryptionMaterial = Arrays.asList(
            mapper.treeToValue(rootNode, RemoteStoreFileEncryptionMaterial[].class));
      }
    }
  }

  private void initPresignedUrls(CommandType commandType, JsonNode jsonNode)
  throws SnowflakeSQLException, JsonProcessingException, IOException
  {
    presignedUrls = new ArrayList<>();
    JsonNode rootNode = jsonNode.path("data").path("presignedUrls");
    if (commandType == CommandType.DOWNLOAD)
    {
      logger.debug("initEncryptionMaterial: DOWNLOAD");

      if (!rootNode.isMissingNode() && !rootNode.isNull())
      {
        presignedUrls = Arrays.asList(mapper.readValue(rootNode.toString(), String[].class));
      }
    }
  }

  public enum CommandType
  {
    UPLOAD,
    DOWNLOAD
  }

  private CommandType commandType = CommandType.UPLOAD;

  private boolean autoCompress = true;

  private boolean overwrite = false;
  private int currentRowIndex;
  private List