package org.apache.tinkerpop.gremlin.orientdb.executor;

import com.orientechnologies.common.exception.OException;
import com.orientechnologies.orient.core.command.script.transformer.OScriptTransformer;
import com.orientechnologies.orient.core.exception.OCommandExecutionException;
import com.orientechnologies.orient.core.sql.executor.OExecutionPlan;
import com.orientechnologies.orient.core.sql.executor.OQueryMetrics;
import com.orientechnologies.orient.core.sql.executor.OResult;
import com.orientechnologies.orient.core.sql.executor.OResultSet;
import java.util.Map;
import java.util.Optional;
import org.apache.tinkerpop.gremlin.orientdb.OrientGraph;
import org.apache.tinkerpop.gremlin.process.traversal.Traversal;

/** Created by Enrico Risa on 24/01/17. */
public class OGremlinScriptResultSet implements OResultSet, OQueryMetrics {

  private final String textTraversal;
  protected Traversal traversal;
  private OScriptTransformer transformer;
  private boolean closeGraph;
  private boolean closing = false;

  long startTime;
  long totalExecutionTime = 0;

  public OGremlinScriptResultSet(
      String iText, Traversal traversal, OScriptTransformer transformer) {
    this(iText, traversal, transformer, false);
  }

  public OGremlinScriptResultSet(
      String iText, Traversal traversal, OScriptTransformer transformer, boolean closeGraph) {
    this.textTraversal = iText;
    this.traversal = traversal;
    this.transformer = transformer;
    this.closeGraph = closeGraph;
    this.startTime = System.currentTimeMillis();
  }

  @Override
  public boolean hasNext() {
    return traversal.hasNext();
  }

  @Override
  public OResult next() {

    long begin = System.currentTimeMillis();

    try {
      Object next = traversal.next();
      return transformer.toResult(next);
    } finally {
      totalExecutionTime += (System.currentTimeMillis() - begin);
    }
  }

  @Override
  public void close() {
    try {
      traversal.close();
      if (closeGraph) {
        traversal
            .asAdmin()
            .getGraph()
            .ifPresent(
                graph -> {
                  try {
                    OrientGraph g = (OrientGraph) graph;
                    if (!closing) {
                      closing = true;
                      g.close();
                    }
                  } catch (Exception e) {
                    throw OException.wrapException(
                        new OCommandExecutionException("Error closing the Graph "), e);
                  }
                });
      }
    } catch (Exception e) {
      throw OException.wrapException(
          new OCommandExecutionException("Error closing the gremlin Result Set"), e);
    }
  }

  @Override
  public Optional<OExecutionPlan> getExecutionPlan() {
    return Optional.empty();
  }

  @Override
  public Map<String, Long> getQueryStats() {
    return null;
  }

  @Override
  public String getStatement() {
    return textTraversal;
  }

  @Override
  public long getStartTime() {
    return startTime;
  }

  @Override
  public long getElapsedTimeMillis() {
    return totalExecutionTime;
  }

  @Override
  public String getLanguage() {
    return "gremlin";
  }
}

@@ -50,16 +50,87 @@
 cutor;%0A%0A
+import com.orientechnologies.common.concur.lock.OInterruptedException;%0A
 import c
@@ -338,32 +338,102 @@
 utionException;%0A
+import com.orientechnologies.orient.core.exception.OStorageException;%0A
 import com.orien
@@ -684,16 +684,65 @@
 ultSet;%0A
+import java.nio.channels.ClosedChannelException;%0A
 import j
@@ -904,16 +904,106 @@
 aversal;
+%0Aimport org.apache.tinkerpop.gremlin.process.traversal.util.TraversalInterruptedException;
 %0A%0A/** Cr
@@ -1861,24 +1861,36 @@
 hasNext() %7B%0A
+    try %7B%0A  
     return t
@@ -1915,183 +1915,901 @@
 ;%0A  
-%7D%0A%0A  @Override%0A  public OResult next() %7B%0A%0A    long begin = System.currentTimeMillis();%0A%0A    try %7B%0A      Object next = traversal.next();%0A      return transformer.toResult(next)
+  %7D catch (TraversalInterruptedException e) %7B%0A      throw new OInterruptedException(%22Timeout expired%22);%0A    %7D catch (OStorageException se) %7B%0A      if (se.getCause() instanceof TraversalInterruptedException%0A          %7C%7C se.getCause() instanceof ClosedChannelException) %7B%0A        throw new OInterruptedException(%22Timeout expired%22);%0A      %7D%0A      throw se;%0A    %7D%0A  %7D%0A%0A  @Override%0A  public OResult next() %7B%0A%0A    long begin = System.currentTimeMillis();%0A%0A    try %7B%0A      Object next = traversal.next();%0A      return transformer.toResult(next);%0A    %7D catch (TraversalInterruptedException e) %7B%0A      throw new OInterruptedException(%22Timeout expired%22);%0A    %7D catch (OStorageException se) %7B%0A      if (se.getCause() instanceof TraversalInterruptedException%0A          %7C%7C se.getCause() instanceof ClosedChannelException) %7B%0A        throw new OInterruptedException(%22Timeout expired%22);%0A      %7D%0A      throw se
 ;%0A  

Fix timeout exception management on traversal iterators