Simple Query Cursor support

src/main/java/com/amazon/opendistroforelasticsearch/sql/domain/Select.java

@@ -50,7 +50,7 @@ public class Select extends Query {
     private Having having;
     private List<Order> orderBys = new ArrayList<>();
     private int offset;
-    private int rowCount = 200;
+    private Integer rowCount;
     private boolean containsSubQueries;
     private List<SubQueryExpression> subQueries;
     private boolean selectAll = false;
@@ -59,6 +59,8 @@ public class Select extends Query {
     public boolean isQuery = false;
     public boolean isAggregate = false;
 
+    public static final int DEFAULT_LIMIT = 200;
+
     public Select() {
     }
 
@@ -70,7 +72,7 @@ public void setOffset(int offset) {
         this.offset = offset;
     }
 
-    public void setRowCount(int rowCount) {
+    public void setRowCount(Integer rowCount) {
         this.rowCount = rowCount;
     }
 
@@ -106,7 +108,7 @@ public int getOffset() {
         return offset;
     }
 
-    public int getRowCount() {
+    public Integer getRowCount() {
         return rowCount;
     }
 

src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/cursor/CursorActionRequestRestExecutorFactory.java

@@ -0,0 +1,36 @@
+/*
+ *   Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.executor.cursor;
+
+import com.amazon.opendistroforelasticsearch.sql.executor.Format;
+import org.elasticsearch.rest.RestRequest;
+
+public class CursorActionRequestRestExecutorFactory {
+    //TODO: add javadocs, see RestExecutor
+
+    public static CursorRestExecutor createExecutor(RestRequest request, String cursor, Format format) {
+        if (isCursorCloseRequest(request)) {
+            return new CursorAsyncRestExecutor(new CursorCloseExecutor(cursor));
+        } else {
+            return new CursorAsyncRestExecutor(new CursorResultExecutor(cursor, format));
+        }
+    }
+
+    private static boolean isCursorCloseRequest(final RestRequest request) {
+        return request.path().endsWith("/_sql/close");
+    }
+
+}

src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/cursor/CursorAsyncRestExecutor.java

@@ -0,0 +1,138 @@
+/*
+ *   Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.executor.cursor;
+
+import com.amazon.opendistroforelasticsearch.sql.esdomain.LocalClusterState;
+import com.amazon.opendistroforelasticsearch.sql.metrics.MetricName;
+import com.amazon.opendistroforelasticsearch.sql.metrics.Metrics;
+import com.amazon.opendistroforelasticsearch.sql.query.QueryAction;
+import com.amazon.opendistroforelasticsearch.sql.query.join.BackOffRetryStrategy;
+import com.amazon.opendistroforelasticsearch.sql.utils.LogUtils;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.common.unit.TimeValue;
+import org.elasticsearch.rest.BytesRestResponse;
+import org.elasticsearch.rest.RestChannel;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.threadpool.ThreadPool;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.Map;
+import java.util.function.Predicate;
+
+import static com.amazon.opendistroforelasticsearch.sql.plugin.SqlSettings.QUERY_SLOWLOG;
+
+public class CursorAsyncRestExecutor implements CursorRestExecutor {
+    /**
+     * Custom thread pool name managed by ES
+     */
+    public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker";
+
+    private static final Logger LOG = LogManager.getLogger(CursorAsyncRestExecutor.class);
+
+    private static final Predicate<QueryAction> ALL_ACTION_IS_BLOCKING = anyAction -> true;
+
+    /**
+     * Delegated rest executor to async
+     */
+    private final CursorRestExecutor executor;
+
+    /**
+     * Request type that expect to async to avoid blocking
+     */
+    private final Predicate<QueryAction> isBlocking;
+
+    CursorAsyncRestExecutor(CursorRestExecutor executor) {
+        this(executor, ALL_ACTION_IS_BLOCKING);
+    }
+
+    CursorAsyncRestExecutor(CursorRestExecutor executor, Predicate<QueryAction> isBlocking) {
+        this.executor = executor;
+        this.isBlocking = isBlocking;
+    }
+
+
+    public void execute(Client client, Map<String, String> params, RestChannel channel) throws Exception {
+        LOG.info("executing something inside CursorAsyncRestExecutor execute ");
+        async(client, params, channel);
+    }
+
+    public String execute(Client client, Map<String, String> params) throws Exception {
+        return "string from CursorAsyncRestExecutor execute()";
+    }
+
+    /**
+     * Run given task in thread pool asynchronously
+     */
+    private void async(Client client, Map<String, String> params, RestChannel channel) {
+
+        ThreadPool threadPool = client.threadPool();
+        Runnable runnable = () -> {
+            try {
+                doExecuteWithTimeMeasured(client, params, channel);
+            } catch (IOException e) {
+                Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+                LOG.warn("[{}] [MCB] async task got an IO/SQL exception: {}", LogUtils.getRequestId(),
+                        e.getMessage());
+                e.printStackTrace();
+                channel.sendResponse(new BytesRestResponse(RestStatus.INTERNAL_SERVER_ERROR, e.getMessage()));
+            } catch (IllegalStateException e) {
+                Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+                LOG.warn("[{}] [MCB] async task got a runtime exception: {}", LogUtils.getRequestId(),
+                        e.getMessage());
+                e.printStackTrace();
+                channel.sendResponse(new BytesRestResponse(RestStatus.INSUFFICIENT_STORAGE,
+                        "Memory circuit is broken."));
+            } catch (Throwable t) {
+                Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+                LOG.warn("[{}] [MCB] async task got an unknown throwable: {}", LogUtils.getRequestId(),
+                        t.getMessage());
+                t.printStackTrace();
+                channel.sendResponse(new BytesRestResponse(RestStatus.INTERNAL_SERVER_ERROR,
+                        String.valueOf(t.getMessage())));
+            } finally {
+                BackOffRetryStrategy.releaseMem(executor);
+            }
+        };
+
+        // Preserve context of calling thread to ensure headers of requests are forwarded when running blocking actions
+        threadPool.schedule(
+                threadPool.preserveContext(LogUtils.withCurrentContext(runnable)),
+                new TimeValue(0L),
+                SQL_WORKER_THREAD_POOL_NAME
+        );
+    }
+
+    /**
+     * Time the real execution of Executor and log slow query for troubleshooting
+     */
+    private void doExecuteWithTimeMeasured(Client client,
+                                           Map<String, String> params,
+                                           RestChannel channel) throws Exception {
+        long startTime = System.nanoTime();
+        try {
+            executor.execute(client, params, channel);
+        } finally {
+            Duration elapsed = Duration.ofNanos(System.nanoTime() - startTime);
+            int slowLogThreshold = LocalClusterState.state().getSettingValue(QUERY_SLOWLOG);
+            if (elapsed.getSeconds() >= slowLogThreshold) {
+                LOG.warn("[{}] Slow query: elapsed={} (ms)", LogUtils.getRequestId(), elapsed.toMillis());
+            }
+        }
+    }
+}

src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/cursor/CursorCloseExecutor.java

@@ -0,0 +1,110 @@
+/*
+ *   Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.sql.executor.cursor;
+
+import com.amazon.opendistroforelasticsearch.sql.metrics.MetricName;
+import com.amazon.opendistroforelasticsearch.sql.metrics.Metrics;
+import com.amazon.opendistroforelasticsearch.sql.rewriter.matchtoterm.VerificationException;
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.search.ClearScrollResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.rest.BytesRestResponse;
+import org.elasticsearch.rest.RestChannel;
+import org.json.JSONException;
+import org.json.JSONObject;
+
+import java.util.Base64;
+import java.util.Map;
+
+import static org.elasticsearch.rest.RestStatus.OK;
+
+public class CursorCloseExecutor implements CursorRestExecutor {
+
+    private static final String SUCCEEDED_TRUE = "{\"succeeded\":true}";
+    private static final String SUCCEEDED_FALSE = "{\"succeeded\":false}";
+
+    private String cursorId;
+
+    public CursorCloseExecutor(String cursorId) {
+        this.cursorId = cursorId;
+    }
+
+    public void execute(Client client, Map<String, String> params, RestChannel channel) throws Exception {
+        try {
+            String formattedResponse = execute(client, params);
+            channel.sendResponse(new BytesRestResponse(OK, "application/json; charset=UTF-8", formattedResponse));
+        } catch (IllegalArgumentException | JSONException e) {
+            Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_CUS).increment();
+            e.printStackTrace();
+            channel.sendResponse(new BytesRestResponse(channel, e));
+        } catch (ElasticsearchException e) {
+            int status = (e.status().getStatus());
+            if (status > 399 && status < 500) {
+                Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_CUS).increment();
+            } else if (status > 499) {
+                Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+            }
+            e.printStackTrace();
+            channel.sendResponse(new BytesRestResponse(channel, e));
+        }
+    }
+
+    public String execute(Client client, Map<String, String> params) throws Exception {
+        String decodedCursorContext = new String(Base64.getDecoder().decode(cursorId));
+        JSONObject cursorJson = new JSONObject(decodedCursorContext);
+
+        String type = cursorJson.optString("type", null); // see if it is a good case to use Optionals
+        CursorType cursorType = null;
+
+        if (type != null) {
+            cursorType = CursorType.valueOf(type);
+        }
+
+        if (cursorType!=null) {
+            switch(cursorType) {
+                case DEFAULT:
+                    return handleDefaultCursorCloseRequest(client, cursorJson);
+                case AGGREGATION:
+                    return handleAggregationCursorCloseRequest(client, cursorJson);
+                case JOIN:
+                    return handleJoinCursorCloseRequest(client, cursorJson);
+                default: throw new VerificationException("Unsupported cursor");
+            }
+        }
+
+        throw new VerificationException("Invalid cursor");
+    }
+
+    private String handleDefaultCursorCloseRequest(Client client, JSONObject cursorContext) {
+        String scrollId = cursorContext.getString("scrollId");
+        ClearScrollResponse clearScrollResponse = client.prepareClearScroll().addScrollId(scrollId).get();
+        if (clearScrollResponse.isSucceeded()) {
+            return SUCCEEDED_TRUE;
+        } else {
+            Metrics.getInstance().getNumericalMetric(MetricName.FAILED_REQ_COUNT_SYS).increment();
+            return SUCCEEDED_FALSE;
+        }
+    }
+
+    private String handleAggregationCursorCloseRequest(Client client, JSONObject cursorContext) {
+        return SUCCEEDED_TRUE;
+    }
+
+    private String handleJoinCursorCloseRequest(Client client, JSONObject cursorContext) {
+        return SUCCEEDED_FALSE;
+    }
+
+}

src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/cursor/CursorRestExecutor.java

@@ -0,0 +1,30 @@
+/*
+ *   Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License").
+ *   You may not use this file except in compliance with the License.
+ *   A copy of the License is located at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   or in the "license" file accompanying this file. This file is distributed
+ *   on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *   express or implied. See the License for the specific language governing
+ *   permissions and limitations under the License.
+ */
+
+
+package com.amazon.opendistroforelasticsearch.sql.executor.cursor;
+
+import org.elasticsearch.client.Client;
+import org.elasticsearch.rest.RestChannel;
+
+import java.util.Map;
+
+public interface CursorRestExecutor {
+
+    void execute(Client client, Map<String, String> params, RestChannel channel)
+            throws Exception;
+
+    String execute(Client client, Map<String, String> params) throws Exception;
+}