Set DocumentOperationType in ReplyContext

1 files changed, 12 insertions, 11 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java b/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
index 56849b9db8f..8f699c156b4 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/documentapi/metrics/DocumentOperationType.java
@@ -1,25 +1,26 @@
 // Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.documentapi.metrics;
 
-import com.yahoo.documentapi.messagebus.protocol.DocumentProtocol;
+import com.yahoo.documentapi.messagebus.protocol.PutDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.RemoveDocumentMessage;
+import com.yahoo.documentapi.messagebus.protocol.UpdateDocumentMessage;
 import com.yahoo.messagebus.Message;
 
 /**
  * @author freva
  */
 public enum DocumentOperationType {
-    PUT, REMOVE, UPDATE;
+    PUT, REMOVE, UPDATE, ERROR;
 
     public static DocumentOperationType fromMessage(Message msg) {
-        switch (msg.getType()) {
-            case DocumentProtocol.MESSAGE_PUTDOCUMENT:
-                return PUT;
-            case DocumentProtocol.MESSAGE_UPDATEDOCUMENT:
-                return UPDATE;
-            case DocumentProtocol.MESSAGE_REMOVEDOCUMENT:
-                return REMOVE;
-            default:
-                return null;
+        if (msg instanceof PutDocumentMessage) {
+            return PUT;
+        } else if (msg instanceof RemoveDocumentMessage) {
+            return REMOVE;
+        } else if (msg instanceof UpdateDocumentMessage) {
+            return UPDATE;
+        } else {
+            return ERROR;
         }
     }
 }