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

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

package net.snowflake.client.ingest;


import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.net.MediaType;
import net.snowflake.client.IngestFilesTester;
import net.snowflake.client.core.SFSession;
import org.apache.http.HttpEntity;
import org.apache.http.HttpResponse;
import org.apache.http.StatusLine;
import org.apache.http.client.HttpClient;
import org.apache.http.client.ResponseHandler;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.utils.URIBuilder;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;
import java.nio.file.Path;
import net.snowflake.client.TestConnectionUtil;

/**
 * Created by hyu on 10/18/16.
 */

public class PushTask implements Callable
{
  private static final String JSON_UTF8_CT = MediaType.JSON_UTF_8.toString();
  private ConcurrentHashMap pushedFiles = new ConcurrentHashMap<>();
  private AtomicInteger batchesRemaining;
  private List> batchOfFiles;
  private String jwtToken;
  private int inter_batch_delay = 0;
  private CloseableHttpClient httpClient;
  private String fqPipeName;
  private ObjectMapper objectMapper = new ObjectMapper();

  public PushTask(String jwtToken,
                  List> batchOfFiles,
                  CloseableHttpClient httpClient,
                  String fqPipeName)
  {
    this.batchOfFiles = batchOfFiles;
    this.jwtToken = jwtToken;
    this.httpClient = httpClient;
    this.fqPipeName = fqPipeName;
    this.batchesRemaining = new AtomicInteger(batchOfFiles.size());
  }

  @Override
  public Integer call() throws Exception
  {
    final ThreadLocalRandom rnd = ThreadLocalRandom.current();
    int cnt = 0;
    while(batchesRemaining.decrementAndGet() >=0 )
    {
      int batchIndex = batchOfFiles.size() - batchesRemaining.get() - 1;
      List files = batchOfFiles.get(batchIndex);
      final Boolean success = postFiles(files, jwtToken);
      if (success) {
        cnt += files.size();
        for (Path file : files)
        {
          final String f = file.getFileName().toString();
          pushedFiles.put(f, f);
        }
      } else {
        throw new IllegalStateException("post failed");
      }
      final int sl = inter_batch_delay +
          rnd.nextInt(1 + inter_batch_delay / 10);
      Thread.sleep(sl);
    }
    //and random amout up to 10% to desync threads...
    return cnt;
  }

  private Boolean postFiles(Collection csv_files, String jwt_token)
      throws Exception
  {
    final URI insert_endpoint = insertFilesEndpoint();
    say("pushing " + csv_files.size() + " files to " + insert_endpoint);
    final HttpPost post = new HttpPost(insert_endpoint);

    addAuthnHeader(post, jwt_token);

    final String body = createPostJsonBody(csv_files);

    final StringEntity entity = new StringEntity(body, StandardCharsets.UTF_8);
    entity.setContentType(JSON_UTF8_CT);
    post.setEntity(entity);

    final ResponseHandler handler = new ResponseHandler()
    {
      @Override
      public Boolean handleResponse(HttpResponse response)
          throws IOException
      {
        final StatusLine statusLine = response.getStatusLine();
        final int statusCode = statusLine.getStatusCode();
        final HttpEntity responseEntity = response.getEntity();
        final String str = EntityUtils.toString(responseEntity);
        if (statusCode == 200) {
          say("insert returned: " + str);
          return true;
        } else {
          say("postFiles request failed got: " + statusLine + "\n:" + str);
          return false;
        }
      }
    };

    final int attempts = 60;

    for (int i = 1; i <= attempts; i++) {
      try {
        final Boolean ok = httpClient.execute(post, handler);
        assert ok != null;
        if (ok) {
          return true; //else retry
        }
      } catch (Exception e) {
        say("got error " + e +
            " during file push on attempt " + i + " of " + attempts);

      }

      Thread.sleep(1000);
    }

    throw new IllegalStateException("giving up on post");
  }

  private URI insertFilesEndpoint() throws URISyntaxException
  {
    return IngestFilesTester.insertFilesEndpoint(fqPipeName);
  }

  private static void say(String arg)
  {
    System.out.println(System.currentTimeMillis() + ":" +
        Thread.currentThread().getId() + "  " + arg);
  }

  public void setInterBatchDelay(int inter_batch_delay)
  {
    this.inter_batch_delay = inter_batch_delay;
  }

  private String createPostJsonBody(Collection csv_files)
      throws IOException
  {
    IngestPushFilesRequest req = new IngestPushFilesRequest();
    final List files =
        new ArrayList<>(csv_files.size());

    for (Path csv_file : csv_files)
    {
      final IngestPushFilesRequest.IngestFile file =
          new IngestPushFilesRequest.IngestFile();
      file.path = csv_file.getFileName().toString();
      try
      {
        file.size = Files.size(csv_file);
      }
      catch (IOException ex)
      {
        file.size = 100L;
      }
      files.add(file);
    }

    req.files = files;

    return this.objectMapper.writeValueAsString(req);
  }

  public static class IngestPushFilesRequest {
    public List files;

    public static final class IngestFile
    {
      public String path;
      public Long size;  //optional
    }
  }

  public ConcurrentHashMap getPushedFiles()
  {
    return this.pushedFiles;
  }

  private static void addAuthnHeader(HttpUriRequest request,
                                    String jwt_token)
  {
    request.setHeader(SFSession.SF_HEADER_AUTHORIZATION,
        "Bearer " + jwt_token);
  }
}