Add PageSink API to rollback writes on failures

presto-base-jdbc/src/main/java/com/facebook/presto/plugin/jdbc/JdbcRecordSink.java

@@ -171,6 +171,20 @@ public String commit()
  return ""; // the committer does not need any additional info
  }
 
+ @SuppressWarnings("UnusedDeclaration")
+ @Override
+ public void rollback()
+ {
+ // rollback and close
+ try (Connection connection = this.connection;
+ PreparedStatement statement = this.statement) {
+ connection.rollback();
+ }
+ catch (SQLException e) {
+ throw Throwables.propagate(e);
+ }
+ }
+
  @Override
  public List<Type> getColumnTypes()
  {

presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraRecordSink.java

@@ -149,6 +149,9 @@ public String commit()
  return "";
  }
 
+ @Override
+ public void rollback() {}
+
  @Override
  public List<Type> getColumnTypes()
  {

presto-hive/src/main/java/com/facebook/presto/hive/HiveRecordSink.java

@@ -205,6 +205,17 @@ public String commit()
  return ""; // the committer can list the directory
  }
 
+ @Override
+ public void rollback()
+ {
+ try {
+ recordWriter.close(true);
+ }
+ catch (IOException e) {
+ throw Throwables.propagate(e);
+ }
+ }
+
  @Override
  public List<Type> getColumnTypes()
  {

presto-main/src/main/java/com/facebook/presto/operator/TableWriterOperator.java

@@ -103,6 +103,8 @@ private enum State
 
  private State state = State.RUNNING;
  private long rowCount;
+ private boolean committed;
+ private boolean closed;
 
  public TableWriterOperator(OperatorContext operatorContext,
  ConnectorPageSink pageSink,
@@ -174,11 +176,24 @@ public Page getOutput()
  state = State.FINISHED;
 
  String fragment = pageSink.commit();
+ committed = true;
 
  PageBuilder page = new PageBuilder(TYPES);
  page.declarePosition();
  BIGINT.writeLong(page.getBlockBuilder(0), rowCount);
  VARCHAR.writeSlice(page.getBlockBuilder(1), Slices.utf8Slice(fragment));
  return page.build();
  }
+
+ @Override
+ public void close()
+ throws Exception
+ {
+ if (!closed) {
+ closed = true;
+ if (!committed) {
+ pageSink.rollback();
+ }
+ }
+ }
 }

presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorPageSink.java

@@ -109,6 +109,12 @@ public String commit()
  return Joiner.on(':').join(nodeId, Joiner.on(",").join(shardUuids));
  }
 
+ @Override
+ public void rollback()
+ {
+ // TODO: clean up open resources
+ }
+
  /**
  * @return page with the sampleWeightBlock at the sampleWeightField index
  */

presto-spi/src/main/java/com/facebook/presto/spi/ConnectorPageSink.java

@@ -20,4 +20,6 @@ public interface ConnectorPageSink
  void appendPage(Page page, Block sampleWeightBlock);
 
  String commit();
+
+ void rollback();
 }

presto-spi/src/main/java/com/facebook/presto/spi/RecordPageSink.java

@@ -38,6 +38,12 @@ public String commit()
  return recordSink.commit();
  }
 
+ @Override
+ public void rollback()
+ {
+ recordSink.rollback();
+ }
+
  @Override
  public void appendPage(Page page, Block sampleWeightBlock)
  {

presto-spi/src/main/java/com/facebook/presto/spi/RecordSink.java

@@ -39,5 +39,7 @@ public interface RecordSink
 
  String commit();
 
+ void rollback();
+
  List<Type> getColumnTypes();
 }