Fix plugin to work against NoteDb

Change-Id: I53f8d54ed02864db34a3a96b130eedf82888c688
diff --git a/src/main/java/com/googlesource/gerrit/plugins/wip/BaseAction.java b/src/main/java/com/googlesource/gerrit/plugins/wip/BaseAction.java
index ae25811..f390369 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/wip/BaseAction.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/wip/BaseAction.java
@@ -20,6 +20,7 @@
 import com.google.gerrit.reviewdb.client.Change;
 import com.google.gerrit.reviewdb.client.Change.Status;
 import com.google.gerrit.reviewdb.client.ChangeMessage;
+import com.google.gerrit.reviewdb.client.PatchSet;
 import com.google.gerrit.reviewdb.server.ReviewDb;
 import com.google.gerrit.server.ChangeUtil;
 import com.google.gerrit.server.CurrentUser;
@@ -27,6 +28,7 @@
 import com.google.gerrit.server.git.BatchUpdate;
 import com.google.gerrit.server.git.UpdateException;
 import com.google.gerrit.server.index.change.ChangeIndexer;
+import com.google.gerrit.server.notedb.ChangeUpdate;
 import com.google.gwtorm.server.OrmException;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
@@ -55,9 +57,9 @@
     this.batchUpdateFactory = batchUpdateFactory;
   }
 
-  protected void changeStatus(Change change, Input input, final Status from,
-      final Status to) throws OrmException, RestApiException,
-      IOException, UpdateException {
+  protected void changeStatus(Change change, final PatchSet.Id psId,
+      Input input, final Status from, final Status to)
+      throws OrmException, RestApiException, IOException, UpdateException {
     ReviewDb db = dbProvider.get();
     Change.Id changeId = change.getId();
     db.changes().beginTransaction(changeId);
@@ -67,8 +69,10 @@
         @Override
         public boolean updateChange(BatchUpdate.ChangeContext ctx) {
           Change change = ctx.getChange();
+          ChangeUpdate update = ctx.getUpdate(psId);
           if (change.getStatus() == from) {
             change.setStatus(to);
+            update.setStatus(change.getStatus());
             return true;
           }
           return false;
diff --git a/src/main/java/com/googlesource/gerrit/plugins/wip/ReadyForReview.java b/src/main/java/com/googlesource/gerrit/plugins/wip/ReadyForReview.java
index f022867..9f852eb 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/wip/ReadyForReview.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/wip/ReadyForReview.java
@@ -58,7 +58,8 @@
       throw new ResourceConflictException("not current patch set");
     }
 
-    changeStatus(change, input, Status.DRAFT, Status.NEW);
+    changeStatus(change, rsrc.getPatchSet().getId(), input, Status.DRAFT,
+        Status.NEW);
     return Response.none();
   }
 
diff --git a/src/main/java/com/googlesource/gerrit/plugins/wip/WorkInProgress.java b/src/main/java/com/googlesource/gerrit/plugins/wip/WorkInProgress.java
index b84fe1e..05a3ebf 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/wip/WorkInProgress.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/wip/WorkInProgress.java
@@ -58,7 +58,8 @@
       throw new ResourceConflictException("not current patch set");
     }
 
-    changeStatus(change, input, Status.NEW, Status.DRAFT);
+    changeStatus(change, rsrc.getPatchSet().getId(), input, Status.NEW,
+        Status.DRAFT);
     return Response.none();
   }