Add thrift serde support for TaskStatus

Added Drift annotations to TaskStatus and related classes to support thrift serde.
prestodb · Oct 14, 2020 · 118b48a · 118b48a
1 parent 213660b
commit 118b48a
Show file tree

Hide file tree

Showing 13 changed files with 340 additions and 16 deletions.
diff --git a/pom.xml b/pom.xml
@@ -45,15 +45,15 @@
         <air.maven.version>3.3.9</air.maven.version>
 
         <dep.antlr.version>4.7.1</dep.antlr.version>
-        <dep.airlift.version>0.194</dep.airlift.version>
+        <dep.airlift.version>0.195</dep.airlift.version>
         <dep.packaging.version>${dep.airlift.version}</dep.packaging.version>
         <dep.slice.version>0.38</dep.slice.version>
         <dep.testing-mysql-server-5.version>0.6</dep.testing-mysql-server-5.version>
         <dep.aws-sdk.version>1.11.697</dep.aws-sdk.version>
         <dep.okhttp.version>3.9.0</dep.okhttp.version>
         <dep.jdbi3.version>3.4.0</dep.jdbi3.version>
         <dep.oracle.version>19.3.0.0</dep.oracle.version>
-        <dep.drift.version>1.28</dep.drift.version>
+        <dep.drift.version>1.29</dep.drift.version>
         <dep.joda.version>2.10.5</dep.joda.version>
         <dep.tempto.version>1.50</dep.tempto.version>
         <dep.testng.version>6.10</dep.testng.version>

diff --git a/presto-client/pom.xml b/presto-client/pom.xml
@@ -56,6 +56,11 @@
             <artifactId>security</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.facebook.drift</groupId>
+            <artifactId>drift-api</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>io.airlift</groupId>
             <artifactId>units</artifactId>

diff --git a/presto-client/src/main/java/com/facebook/presto/client/ErrorLocation.java b/presto-client/src/main/java/com/facebook/presto/client/ErrorLocation.java
@@ -13,6 +13,9 @@
  */
 package com.facebook.presto.client;
 
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
@@ -22,12 +25,14 @@
 import static com.google.common.base.Preconditions.checkArgument;
 
 @Immutable
+@ThriftStruct
 public class ErrorLocation
 {
     private final int lineNumber;
     private final int columnNumber;
 
     @JsonCreator
+    @ThriftConstructor
     public ErrorLocation(
             @JsonProperty("lineNumber") int lineNumber,
             @JsonProperty("columnNumber") int columnNumber)
@@ -40,12 +45,14 @@ public ErrorLocation(
     }
 
     @JsonProperty
+    @ThriftField(1)
     public int getLineNumber()
     {
         return lineNumber;
     }
 
     @JsonProperty
+    @ThriftField(2)
     public int getColumnNumber()
     {
         return columnNumber;

diff --git a/presto-main/pom.xml b/presto-main/pom.xml
@@ -206,6 +206,11 @@
             <artifactId>drift-client</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>com.facebook.drift</groupId>
+            <artifactId>drift-protocol</artifactId>
+        </dependency>
+
         <dependency>
             <groupId>com.facebook.drift</groupId>
             <artifactId>drift-codec</artifactId>

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/ExecutionFailureInfo.java b/presto-main/src/main/java/com/facebook/presto/execution/ExecutionFailureInfo.java
@@ -13,6 +13,9 @@
  */
 package com.facebook.presto.execution;
 
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
 import com.facebook.presto.client.ErrorLocation;
 import com.facebook.presto.client.FailureInfo;
 import com.facebook.presto.spi.ErrorCode;
@@ -28,10 +31,13 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
+import static com.facebook.drift.annotations.ThriftField.Recursiveness.TRUE;
+import static com.facebook.drift.annotations.ThriftField.Requiredness.OPTIONAL;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.Objects.requireNonNull;
 
 @Immutable
+@ThriftStruct
 public class ExecutionFailureInfo
 {
     private static final Pattern STACK_TRACE_PATTERN = Pattern.compile("(.*)\\.(.*)\\(([^:]*)(?::(.*))?\\)");
@@ -47,6 +53,7 @@ public class ExecutionFailureInfo
     private final HostAddress remoteHost;
 
     @JsonCreator
+    @ThriftConstructor
     public ExecutionFailureInfo(
             @JsonProperty("type") String type,
             @JsonProperty("message") String message,
@@ -72,53 +79,61 @@ public ExecutionFailureInfo(
     }
 
     @JsonProperty
+    @ThriftField(1)
     public String getType()
     {
         return type;
     }
 
     @Nullable
     @JsonProperty
+    @ThriftField(2)
     public String getMessage()
     {
         return message;
     }
 
     @Nullable
     @JsonProperty
+    @ThriftField(value = 3, isRecursive = TRUE, requiredness = OPTIONAL)
     public ExecutionFailureInfo getCause()
     {
         return cause;
     }
 
     @JsonProperty
+    @ThriftField(4)
     public List<ExecutionFailureInfo> getSuppressed()
     {
         return suppressed;
     }
 
     @JsonProperty
+    @ThriftField(5)
     public List<String> getStack()
     {
         return stack;
     }
 
     @Nullable
     @JsonProperty
+    @ThriftField(6)
     public ErrorLocation getErrorLocation()
     {
         return errorLocation;
     }
 
     @Nullable
     @JsonProperty
+    @ThriftField(7)
     public ErrorCode getErrorCode()
     {
         return errorCode;
     }
 
     @Nullable
     @JsonProperty
+    @ThriftField(8)
     public HostAddress getRemoteHost()
     {
         return remoteHost;

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/Lifespan.java b/presto-main/src/main/java/com/facebook/presto/execution/Lifespan.java
@@ -14,6 +14,9 @@
 
 package com.facebook.presto.execution;
 
+import com.facebook.drift.annotations.ThriftConstructor;
+import com.facebook.drift.annotations.ThriftField;
+import com.facebook.drift.annotations.ThriftStruct;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
 
@@ -23,6 +26,7 @@
 import static com.google.common.base.Preconditions.checkState;
 import static java.lang.Integer.parseInt;
 
+@ThriftStruct
 public class Lifespan
 {
     private static final Lifespan TASK_WIDE = new Lifespan(false, 0);
@@ -40,8 +44,10 @@ public static Lifespan driverGroup(int id)
         return new Lifespan(true, id);
     }
 
-    private Lifespan(boolean grouped, int groupId)
+    @ThriftConstructor
+    public Lifespan(boolean grouped, int groupId)
     {
+        checkArgument(grouped || groupId == 0);
         this.grouped = grouped;
         this.groupId = groupId;
     }
@@ -51,9 +57,16 @@ public boolean isTaskWide()
         return !grouped;
     }
 
+    @ThriftField(1)
+    public boolean isGrouped()
+    {
+        return grouped;
+    }
+
+    @ThriftField(value = 2, name = "groupId")
     public int getId()
     {
-        checkState(grouped);
+        checkState(grouped || groupId == 0);
         return groupId;
     }
 

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/TaskState.java b/presto-main/src/main/java/com/facebook/presto/execution/TaskState.java
@@ -13,47 +13,53 @@
  */
 package com.facebook.presto.execution;
 
+import com.facebook.drift.annotations.ThriftEnum;
+import com.facebook.drift.annotations.ThriftEnumValue;
+
 import java.util.Set;
 import java.util.stream.Stream;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
+@ThriftEnum
 public enum TaskState
 {
     /**
      * Task is planned but has not been scheduled yet.  A task will
      * be in the planned state until, the dependencies of the task
      * have begun producing output.
      */
-    PLANNED(false),
+    PLANNED(0, false),
     /**
      * Task is running.
      */
-    RUNNING(false),
+    RUNNING(1, false),
     /**
      * Task has finished executing and all output has been consumed.
      */
-    FINISHED(true),
+    FINISHED(2, true),
     /**
      * Task was canceled by a user.
      */
-    CANCELED(true),
+    CANCELED(3, true),
     /**
      * Task was aborted due to a failure in the query.  The failure
      * was not in this task.
      */
-    ABORTED(true),
+    ABORTED(4, true),
     /**
      * Task execution failed.
      */
-    FAILED(true);
+    FAILED(5, true);
 
     public static final Set<TaskState> TERMINAL_TASK_STATES = Stream.of(TaskState.values()).filter(TaskState::isDone).collect(toImmutableSet());
 
+    private final int code;
     private final boolean doneState;
 
-    TaskState(boolean doneState)
+    TaskState(int code, boolean doneState)
     {
+        this.code = code;
         this.doneState = doneState;
     }
 
@@ -64,4 +70,10 @@ public boolean isDone()
     {
         return doneState;
     }
+
+    @ThriftEnumValue
+    public int getCode()
+    {
+        return code;
+    }
 }