embulk-input-marketo 0.5.7.alpha.6 → 0.6.0.alpha.1

data/src/main/java/org/embulk/input/marketo/delegate/MarketoBaseInputPluginDelegate.java

@@ -5,17 +5,23 @@ import org.embulk.base.restclient.DefaultServiceDataSplitter;
 import org.embulk.base.restclient.RestClientInputPluginDelegate;
 import org.embulk.base.restclient.RestClientInputTaskBase;
 import org.embulk.base.restclient.ServiceDataSplitter;
+import org.embulk.base.restclient.record.RecordImporter;
+import org.embulk.base.restclient.record.ServiceRecord;
 import org.embulk.config.Config;
 import org.embulk.config.ConfigDefault;
 import org.embulk.config.ConfigDiff;
 import org.embulk.config.TaskReport;
+import org.embulk.input.marketo.MarketoService;
+import org.embulk.input.marketo.MarketoServiceImpl;
 import org.embulk.input.marketo.rest.MarketoRestClient;
 import org.embulk.spi.Exec;
+import org.embulk.spi.PageBuilder;
 import org.embulk.spi.Schema;
 import org.embulk.util.retryhelper.jetty92.DefaultJetty92ClientCreator;
 import org.embulk.util.retryhelper.jetty92.Jetty92RetryHelper;
 import org.joda.time.DateTime;
 
+import java.util.Iterator;
 import java.util.List;
 
 /**
@@ -68,6 +74,22 @@ public abstract class MarketoBaseInputPluginDelegate<T extends MarketoBaseInputP
         task.setJobStartTime(DateTime.now());
     }
 
+    @Override
+    public TaskReport ingestServiceData(T task, RecordImporter recordImporter, int taskIndex, PageBuilder pageBuilder)
+    {
+        MarketoService marketoService = new MarketoServiceImpl(createMarketoRestClient(task));
+        Iterator<ServiceRecord> serviceRecords = getServiceRecords(marketoService, task);
+        int imported = 0;
+        while (serviceRecords.hasNext() && (imported < PREVIEW_RECORD_LIMIT || !Exec.isPreview())) {
+            ServiceRecord next = serviceRecords.next();
+            recordImporter.importRecord(next, pageBuilder);
+            imported++;
+        }
+        return Exec.newTaskReport();
+    }
+
+    protected abstract Iterator<ServiceRecord> getServiceRecords(MarketoService marketoService, T task);
+
     @VisibleForTesting
     public MarketoRestClient createMarketoRestClient(PluginTask task)
     {

data/src/main/java/org/embulk/input/marketo/model/MarketoBulkExtractRequest.java

@@ -2,6 +2,7 @@ package org.embulk.input.marketo.model;
 
 import org.embulk.input.marketo.model.filter.MarketoFilter;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -15,7 +16,7 @@ public class MarketoBulkExtractRequest
 
     private Map<String, String> columnHeaderNames;
 
-    private Map<String, MarketoFilter> filter;
+    private Map<String, MarketoFilter> filter = new HashMap<>();
 
     public List<String> getFields()
     {
@@ -56,4 +57,14 @@ public class MarketoBulkExtractRequest
     {
         this.filter = filter;
     }
+
+    @Override
+    public String toString()
+    {
+        return "MarketoBulkExtractRequest{" +
+                "format='" + format + '\'' +
+                ", columnHeaderNames=" + columnHeaderNames +
+                ", filter=" + filter +
+                '}';
+    }
 }

data/src/main/java/org/embulk/input/marketo/model/MarketoField.java

@@ -15,13 +15,15 @@ public class MarketoField
 
     private MarketoDataType marketoDataType;
 
+    public MarketoField(){}
 
     public MarketoField(String name, String dataType)
     {
         this.name = name;
         try {
             marketoDataType = MarketoDataType.valueOf(dataType.toUpperCase());
-        } catch (IllegalArgumentException ex) {
+        }
+        catch (IllegalArgumentException ex) {
             marketoDataType = MarketoDataType.STRING;
         }
     }
@@ -30,7 +32,6 @@ public class MarketoField
     {
         this.name = name;
         this.marketoDataType = marketoDataType;
-
     }
 
     public String getName()
@@ -43,6 +44,31 @@ public class MarketoField
         return marketoDataType;
     }
 
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+
+        MarketoField field = (MarketoField) o;
+
+        if (name != null ? !name.equals(field.name) : field.name != null) {
+            return false;
+        }
+        return marketoDataType == field.marketoDataType;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        int result = name != null ? name.hashCode() : 0;
+        result = 31 * result + (marketoDataType != null ? marketoDataType.hashCode() : 0);
+        return result;
+    }
 
     public enum MarketoDataType
     {

data/src/main/java/org/embulk/input/marketo/model/MarketoResponse.java

@@ -1,5 +1,6 @@
 package org.embulk.input.marketo.model;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -17,7 +18,7 @@ public class MarketoResponse<T>
 
     private List<MarketoError> errors;
 
-    private List<T> result;
+    private List<T> result = new ArrayList<>();
 
     public String getRequestId()
     {

data/src/main/java/org/embulk/input/marketo/model/filter/DateRangeFilter.java

@@ -28,4 +28,13 @@ public class DateRangeFilter implements MarketoFilter
     {
         this.endAt = endAt;
     }
+
+    @Override
+    public String toString()
+    {
+        return "DateRangeFilter{" +
+                "startAt='" + startAt + '\'' +
+                ", endAt='" + endAt + '\'' +
+                '}';
+    }
 }

data/src/main/java/org/embulk/input/marketo/rest/MarketoBaseRestClient.java

@@ -21,6 +21,7 @@ import org.embulk.util.retryhelper.jetty92.Jetty92SingleRequester;
 import org.embulk.util.retryhelper.jetty92.StringJetty92ResponseEntityReader;
 import org.slf4j.Logger;
 
+import java.io.EOFException;
 import java.io.IOException;
 import java.net.SocketTimeoutException;
 import java.nio.charset.StandardCharsets;
@@ -186,6 +187,9 @@ public class MarketoBaseRestClient implements AutoCloseable
                 if (exception instanceof ExecutionException) {
                     this.toRetry((Exception) exception.getCause());
                 }
+                if (exception instanceof EOFException) {
+                    return true;
+                }
                 if (exception instanceof MarketoAPIException) {
                     //Retry Authenticate Exception
                     MarketoError error = ((MarketoAPIException) exception).getMarketoErrors().get(0);

data/src/main/java/org/embulk/input/marketo/rest/{MarketoFileResponseEntityReader.java → MarketoInputStreamResponseEntityReader.java}

@@ -19,7 +19,7 @@ import java.util.concurrent.TimeUnit;
 /**
  * Created by tai.khuu on 9/5/17.
  */
-public class MarketoFileResponseEntityReader implements Jetty92ResponseReader<InputStream>
+public class MarketoInputStreamResponseEntityReader implements Jetty92ResponseReader<InputStream>
 {
     private static final ObjectReader OBJECT_READER = new ObjectMapper().readerFor(new TypeReference<MarketoResponse<ObjectNode>>(){ });
 
@@ -27,7 +27,7 @@ public class MarketoFileResponseEntityReader implements Jetty92ResponseReader<In
 
     private long timeout;
 
-    public MarketoFileResponseEntityReader(long timeout)
+    public MarketoInputStreamResponseEntityReader(long timeout)
     {
         this.listener = new InputStreamResponseListener();
         this.timeout = timeout;

data/src/main/java/org/embulk/input/marketo/rest/MarketoResponseJetty92EntityReader.java

@@ -82,7 +82,10 @@ public class MarketoResponseJetty92EntityReader<T> implements Jetty92ResponseRea
     {
         InputStream inputStream = this.listener.getInputStream();
         try (InputStreamReader inputStreamReader = new InputStreamReader(inputStream, StandardCharsets.UTF_8)) {
-            return CharStreams.toString(inputStreamReader);
+            String reponseContent = CharStreams.toString(inputStreamReader);
+            //Reset this listener
+            this.listener = new InputStreamResponseListener();
+            return reponseContent;
         }
     }
 }

data/src/main/java/org/embulk/input/marketo/rest/MarketoRestClient.java

@@ -3,6 +3,7 @@ package org.embulk.input.marketo.rest;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Multimap;
 import org.apache.commons.lang3.StringUtils;
@@ -15,7 +16,6 @@ import org.embulk.input.marketo.model.MarketoError;
 import org.embulk.input.marketo.model.MarketoField;
 import org.embulk.input.marketo.model.MarketoResponse;
 import org.embulk.input.marketo.model.filter.DateRangeFilter;
-import org.embulk.input.marketo.model.filter.ListFilter;
 import org.embulk.input.marketo.model.filter.MarketoFilter;
 import org.embulk.spi.DataException;
 import org.embulk.spi.Exec;
@@ -41,10 +41,16 @@ public class MarketoRestClient extends MarketoBaseRestClient
 
     private static final String NEXT_PAGE_TOKEN = "nextPageToken";
 
+    private static final String OFFSET = "offset";
+
+    private static final String MAX_RETURN = "maxReturn";
+
     private String endPoint;
 
     private Integer batchSize;
 
+    private Integer maxReturn;
+
     private static final Logger LOGGER = Exec.getLogger(MarketoRestClient.class.getCanonicalName());
 
     private static final Map<String, Type> TYPE_MAPPING = new ImmutableMap.Builder<String, Type>()
@@ -85,18 +91,24 @@ public class MarketoRestClient extends MarketoBaseRestClient
         @ConfigDefault("300")
         Integer getBatchSize();
         void setBatchSize(Integer batchSize);
+
+        @Config("max_return")
+        @ConfigDefault("200")
+        Integer getMaxReturn();
+        void setMaxReturn(Integer maxReturn);
     }
 
     public MarketoRestClient(PluginTask task, Jetty92RetryHelper retryHelper)
     {
-        this("https://" + task.getAccountId() + ".mktorest.com", "https://" + task.getAccountId() + ".mktorest.com/identity", task.getClientId(), task.getClientSecret(), task.getBatchSize(), task.getMarketoLimitIntervalMilis(), retryHelper);
+        this(MarketoUtils.getEndPoint(task.getAccountId()), MarketoUtils.getIdentityEndPoint(task.getAccountId()), task.getClientId(), task.getClientSecret(), task.getBatchSize(), task.getMaxReturn(), task.getMarketoLimitIntervalMilis(), retryHelper);
     }
 
-    public MarketoRestClient(String endPoint, String identityEndPoint, String clientId, String clientSecret, Integer batchSize, int marketoLimitIntervalMilis, Jetty92RetryHelper retryHelper)
+    public MarketoRestClient(String endPoint, String identityEndPoint, String clientId, String clientSecret, Integer batchSize, Integer maxReturn, int marketoLimitIntervalMilis, Jetty92RetryHelper retryHelper)
     {
         super(identityEndPoint, clientId, clientSecret, marketoLimitIntervalMilis, retryHelper);
         this.endPoint = endPoint;
         this.batchSize = batchSize;
+        this.maxReturn = maxReturn;
     }
 
     public List<MarketoField> describeLead()
@@ -119,37 +131,32 @@ public class MarketoRestClient extends MarketoBaseRestClient
         return TYPE_MAPPING.containsKey(dataType.toLowerCase()) ? TYPE_MAPPING.get(dataType.toLowerCase()) : Types.STRING;
     }
 
-    public String createLeadBulkExtract(Date startTime, Date endTime, List<String> extractFields)
+    public String createLeadBulkExtract(Date startTime, Date endTime, List<String> extractFields, String fitlerField)
     {
-        SimpleDateFormat timeFormat = new SimpleDateFormat(MarketoUtils.MARKETO_DATE_SIMPLE_DATE_FORMAT);
-        MarketoBulkExtractRequest marketoBulkExtractRequest = new MarketoBulkExtractRequest();
-        marketoBulkExtractRequest.setFields(extractFields);
-        marketoBulkExtractRequest.setFormat("CSV");
-        Map<String, MarketoFilter> filterMap = new HashMap<>();
-        DateRangeFilter dateRangeFilter = new DateRangeFilter();
-        dateRangeFilter.setStartAt(timeFormat.format(startTime));
-        dateRangeFilter.setEndAt(timeFormat.format(endTime));
-        filterMap.put("createdAt", dateRangeFilter);
-        marketoBulkExtractRequest.setFilter(filterMap);
+        MarketoBulkExtractRequest marketoBulkExtractRequest = getMarketoBulkExtractRequest(startTime, endTime, extractFields, fitlerField);
         return sendCreateBulkExtractRequest(marketoBulkExtractRequest, MarketoRESTEndpoint.CREATE_LEAD_EXTRACT);
     }
 
-    public String createActitvityExtract(Date startTime, Date endTime, List<String> activityTypes)
+    private MarketoBulkExtractRequest getMarketoBulkExtractRequest(Date startTime, Date endTime, List<String> extractFields, String rangeFilterName)
     {
         SimpleDateFormat timeFormat = new SimpleDateFormat(MarketoUtils.MARKETO_DATE_SIMPLE_DATE_FORMAT);
         MarketoBulkExtractRequest marketoBulkExtractRequest = new MarketoBulkExtractRequest();
+        if (extractFields != null) {
+            marketoBulkExtractRequest.setFields(extractFields);
+        }
         marketoBulkExtractRequest.setFormat("CSV");
         Map<String, MarketoFilter> filterMap = new HashMap<>();
         DateRangeFilter dateRangeFilter = new DateRangeFilter();
         dateRangeFilter.setStartAt(timeFormat.format(startTime));
         dateRangeFilter.setEndAt(timeFormat.format(endTime));
-        filterMap.put("createdAt", dateRangeFilter);
-        if (activityTypes != null) {
-            ListFilter activitiesTypeFilter = new ListFilter();
-            activitiesTypeFilter.addAll(activityTypes);
-            filterMap.put("activities", activitiesTypeFilter);
-        }
+        filterMap.put(rangeFilterName, dateRangeFilter);
         marketoBulkExtractRequest.setFilter(filterMap);
+        return marketoBulkExtractRequest;
+    }
+
+    public String createActivityExtract(Date startTime, Date endTime)
+    {
+        MarketoBulkExtractRequest marketoBulkExtractRequest = getMarketoBulkExtractRequest(startTime, endTime, null, "createdAt");
         return sendCreateBulkExtractRequest(marketoBulkExtractRequest, MarketoRESTEndpoint.CREATE_ACTIVITY_EXTRACT);
     }
 
@@ -157,6 +164,7 @@ public class MarketoRestClient extends MarketoBaseRestClient
     {
         MarketoResponse<ObjectNode> marketoResponse = null;
         try {
+            LOGGER.info("Send bulk extract request [{}]", request);
             marketoResponse = doPost(endPoint + endpoint.getEndpoint(), null, null, OBJECT_MAPPER.writeValueAsString(request), new MarketoResponseJetty92EntityReader<ObjectNode>(READ_TIMEOUT_MILLIS));
         }
         catch (JsonProcessingException e) {
@@ -265,67 +273,93 @@ public class MarketoRestClient extends MarketoBaseRestClient
 
     private InputStream getBulkExtractResult(MarketoRESTEndpoint endpoint, String exportId)
     {
-        InputStream fileStream = doGet(this.endPoint + endpoint.getEndpoint(new ImmutableMap.Builder().put("export_id", exportId).build()), null, null, new MarketoFileResponseEntityReader(READ_TIMEOUT_MILLIS));
-        return fileStream;
+        return doGet(this.endPoint + endpoint.getEndpoint(new ImmutableMap.Builder().put("export_id", exportId).build()), null, null, new MarketoInputStreamResponseEntityReader(READ_TIMEOUT_MILLIS));
     }
 
     public RecordPagingIterable<ObjectNode> getLists()
     {
-        return getRecordWithPagination(endPoint + MarketoRESTEndpoint.GET_LISTS.getEndpoint(), null, ObjectNode.class);
+        return getRecordWithTokenPagination(endPoint + MarketoRESTEndpoint.GET_LISTS.getEndpoint(), null, ObjectNode.class);
     }
 
     public RecordPagingIterable<ObjectNode> getPrograms()
     {
-        return getRecordWithPagination(endPoint + MarketoRESTEndpoint.GET_PROGRAMS.getEndpoint(), null, ObjectNode.class);
+        return getRecordWithOffsetPagination(endPoint + MarketoRESTEndpoint.GET_PROGRAMS.getEndpoint(), null, ObjectNode.class);
     }
 
     public RecordPagingIterable<ObjectNode> getLeadsByProgram(String programId, List<String> fieldNames)
     {
         Multimap<String, String> multimap = ArrayListMultimap.create();
         multimap.put("fields", StringUtils.join(fieldNames, ","));
-        return getRecordWithPagination(endPoint + MarketoRESTEndpoint.GET_LEADS_BY_PROGRAM.getEndpoint(new ImmutableMap.Builder().put("program_id", programId).build()), multimap, ObjectNode.class);
+        return getRecordWithTokenPagination(endPoint + MarketoRESTEndpoint.GET_LEADS_BY_PROGRAM.getEndpoint(new ImmutableMap.Builder().put("program_id", programId).build()), multimap, ObjectNode.class);
     }
 
     public RecordPagingIterable<ObjectNode> getLeadsByList(String listId, List<String> fieldNames)
     {
         Multimap<String, String> multimap = ArrayListMultimap.create();
         multimap.put("fields", StringUtils.join(fieldNames, ","));
-        return getRecordWithPagination(endPoint + MarketoRESTEndpoint.GET_LEADS_BY_LIST.getEndpoint(new ImmutableMap.Builder().put("list_id", listId).build()), multimap, ObjectNode.class);
+        return getRecordWithTokenPagination(endPoint + MarketoRESTEndpoint.GET_LEADS_BY_LIST.getEndpoint(new ImmutableMap.Builder().put("list_id", listId).build()), multimap, ObjectNode.class);
     }
 
     public RecordPagingIterable<ObjectNode> getCampaign()
     {
-        return getRecordWithPagination(endPoint + MarketoRESTEndpoint.GET_CAMPAIGN.getEndpoint(), null, ObjectNode.class);
+        return getRecordWithTokenPagination(endPoint + MarketoRESTEndpoint.GET_CAMPAIGN.getEndpoint(), null, ObjectNode.class);
     }
-
-    private <T> RecordPagingIterable<T> getRecordWithPagination(final String endPoint, final Multimap<String, String> parameters, final Class<T> recordClass)
+    private <T> RecordPagingIterable<T> getRecordWithOffsetPagination(final String endPoint, final Multimap<String, String> parameters, final Class<T> recordClass)
     {
-        return new RecordPagingIterable<>(new RecordPagingIterable.PagingFunction<RecordPagingIterable.MarketoPage<T>>()
+        return new RecordPagingIterable<>(new RecordPagingIterable.PagingFunction<RecordPagingIterable.OffsetPage<T>>()
         {
             @Override
-            public RecordPagingIterable.MarketoPage<T> getNextPage(RecordPagingIterable.MarketoPage<T> currentPage)
+            public RecordPagingIterable.OffsetPage<T> getNextPage(RecordPagingIterable.OffsetPage<T> currentPage)
             {
-                Multimap<String, String> params = ArrayListMultimap.create();
-                params.put(NEXT_PAGE_TOKEN, currentPage.getNextPageToken());
-                return gettMarketoPage(params);
+                return getOffsetPage(currentPage.getNextOffSet());
             }
 
             @Override
-            public RecordPagingIterable.MarketoPage<T> getFirstPage()
+            public RecordPagingIterable.OffsetPage<T> getFirstPage()
             {
-                return gettMarketoPage(null);
+                return getOffsetPage(0);
             }
-            private RecordPagingIterable.MarketoPage<T> gettMarketoPage(Multimap<String, String> params)
+
+            private RecordPagingIterable.OffsetPage<T> getOffsetPage(int offset)
+            {
+                ImmutableListMultimap.Builder<String, String> params = new ImmutableListMultimap.Builder<>();
+                params.put(OFFSET, String.valueOf(offset));
+                params.put(MAX_RETURN, String.valueOf(maxReturn));
+                if (parameters != null) {
+                    params.putAll(parameters);
+                }
+                MarketoResponse<T> marketoResponse = doGet(endPoint, null, params.build(), new MarketoResponseJetty92EntityReader<>(READ_TIMEOUT_MILLIS, recordClass));
+                return new RecordPagingIterable.OffsetPage<>(marketoResponse.getResult(), offset + marketoResponse.getResult().size(), marketoResponse.getResult().size() == maxReturn);
+            }
+        });
+    }
+    private <T> RecordPagingIterable<T> getRecordWithTokenPagination(final String endPoint, final Multimap<String, String> parameters, final Class<T> recordClass)
+    {
+        return new RecordPagingIterable<>(new RecordPagingIterable.PagingFunction<RecordPagingIterable.TokenPage<T>>()
+        {
+            @Override
+            public RecordPagingIterable.TokenPage<T> getNextPage(RecordPagingIterable.TokenPage<T> currentPage)
+            {
+                return getTokenPage(currentPage);
+            }
+            @Override
+            public RecordPagingIterable.TokenPage<T> getFirstPage()
+            {
+                return getTokenPage(null);
+            }
+
+            private RecordPagingIterable.TokenPage<T> getTokenPage(RecordPagingIterable.TokenPage page)
             {
-                if (params == null) {
-                    params = ArrayListMultimap.create();
+                ImmutableListMultimap.Builder params = new ImmutableListMultimap.Builder<>();
+                if (page != null) {
+                    params.put(NEXT_PAGE_TOKEN, page.getNextPageToken());
                 }
                 params.put(BATCH_SIZE, String.valueOf(batchSize));
                 if (parameters != null) {
                     params.putAll(parameters);
                 }
-                MarketoResponse<T> marketoResponse = doGet(endPoint, null, params, new MarketoResponseJetty92EntityReader<>(READ_TIMEOUT_MILLIS, recordClass));
-                return new RecordPagingIterable.MarketoPage<>(marketoResponse.getResult(), marketoResponse.getNextPageToken(), marketoResponse.isMoreResult());
+                MarketoResponse<T> marketoResponse = doGet(endPoint, null, params.build(), new MarketoResponseJetty92EntityReader<>(READ_TIMEOUT_MILLIS, recordClass));
+                return new RecordPagingIterable.TokenPage<>(marketoResponse.getResult(), marketoResponse.getNextPageToken(), marketoResponse.getNextPageToken() != null);
             }
         });
     }