Add OperationHandle

kyuubi-common/src/main/scala/org/apache/kyuubi/cli/Handle.scala

@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.cli
+
+import org.apache.hive.service.rpc.thrift.TProtocolVersion
+
+trait Handle {
+  def identifier: HandleIdentifier
+  def protocol: TProtocolVersion
+}

kyuubi-common/src/main/scala/org/apache/kyuubi/cli/HandleIdentifier.scala

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.cli
+
+import java.nio.ByteBuffer
+import java.util.{Objects, UUID}
+
+import org.apache.hive.service.rpc.thrift.THandleIdentifier
+
+case class HandleIdentifier(publicId: UUID, secretId: UUID) {
+
+  def toTHandleIdentifier: THandleIdentifier = {
+    val guid = new Array[Byte](16)
+    val pbb = ByteBuffer.wrap(guid)
+    val secret = new Array[Byte](16)
+    val sbb = ByteBuffer.wrap(secret)
+    pbb.putLong(publicId.getMostSignificantBits)
+    pbb.putLong(publicId.getLeastSignificantBits)
+    sbb.putLong(secretId.getMostSignificantBits)
+    sbb.putLong(secretId.getLeastSignificantBits)
+    new THandleIdentifier(ByteBuffer.wrap(guid), ByteBuffer.wrap(secret))
+  }
+
+  override def hashCode(): Int = (Objects.hashCode(publicId) + 31) * 31 + Objects.hashCode(secretId)
+
+  override def equals(obj: Any): Boolean = obj match {
+    case HandleIdentifier(pid, sid) =>
+      Objects.equals(publicId, pid) && Objects.equals(secretId, sid)
+    case _ => false
+  }
+
+  override def toString: String = publicId.toString
+}
+
+object HandleIdentifier {
+  def apply(): HandleIdentifier = {
+    val publicId = UUID.randomUUID()
+    val secretId = UUID.randomUUID()
+    new HandleIdentifier(publicId, secretId)
+  }
+
+  def apply(tHandleId: THandleIdentifier): HandleIdentifier = {
+    val pbb = ByteBuffer.wrap(tHandleId.getGuid)
+    val publicId = new UUID(pbb.getLong, pbb.getLong)
+
+    val sbb = ByteBuffer.wrap(tHandleId.getSecret)
+    val secretId = new UUID(sbb.getLong, sbb.getLong)
+
+    new HandleIdentifier(publicId, secretId)
+  }
+}

kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationHandle.scala

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.operation
+
+import java.util.Objects
+
+import org.apache.hive.service.rpc.thrift.{TOperationHandle, TProtocolVersion}
+
+import org.apache.kyuubi.cli.{Handle, HandleIdentifier}
+import org.apache.kyuubi.operation.OperationType.OperationType
+
+case class OperationHandle(
+    identifier: HandleIdentifier,
+    typ: OperationType,
+    protocol: TProtocolVersion) extends Handle {
+
+  private var _hasResultSet: Boolean = _
+
+  def setHasResultSet(hasResultSet: Boolean): Unit = _hasResultSet = hasResultSet
+
+  def toTOperationHandle: TOperationHandle = {
+    val tOperationHandle = new TOperationHandle
+    tOperationHandle.setOperationId(identifier.toTHandleIdentifier)
+    tOperationHandle.setOperationType(OperationType.toTOperationType(typ))
+    tOperationHandle.setHasResultSet(_hasResultSet)
+    tOperationHandle
+  }
+
+  override def hashCode(): Int = {
+    (Objects.hashCode(identifier) + 31) * 31 + Objects.hashCode(typ)
+  }
+
+  override def equals(obj: Any): Boolean = obj match {
+    case OperationHandle(id, typ, _) =>
+      Objects.equals(this.typ, typ) && Objects.equals(this.identifier, id)
+    case _ => false
+  }
+
+  override def toString: String = {
+    s"OperationHandle [type=$typ, identifier: $identifier]"
+  }
+}
+
+object OperationHandle {
+  def apply(operationType: OperationType, protocol: TProtocolVersion): OperationHandle = {
+    apply(HandleIdentifier(), operationType, protocol)
+  }
+
+  def apply(tOperationHandle: TOperationHandle, protocol: TProtocolVersion): OperationHandle = {
+    apply(HandleIdentifier(tOperationHandle.getOperationId),
+      OperationType.getOperationType(tOperationHandle.getOperationType),
+      protocol)
+  }
+
+  def apply(tOperationHandle: TOperationHandle): OperationHandle = {
+    apply(tOperationHandle, TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V1)
+  }
+}

kyuubi-common/src/main/scala/org/apache/kyuubi/operation/OperationStatus.scala

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.operation
+
+import org.apache.kyuubi.KyuubiSQLException
+import org.apache.kyuubi.operation.OperationState.OperationState
+
+case class OperationStatus(
+    state: OperationState,
+    start: Long,
+    completed: Long,
+    hasResultSet: Boolean,
+    exception: Option[KyuubiSQLException] = None)

kyuubi-common/src/test/scala/org/apache/kyuubi/cli/HandleIdentifierSuite.scala

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.cli
+
+import org.apache.kyuubi.KyuubiFunSuite
+
+class HandleIdentifierSuite extends KyuubiFunSuite {
+
+  test("HandleIdentifier") {
+    val id1 = HandleIdentifier()
+    val tid1 = id1.toTHandleIdentifier
+    val id2 = HandleIdentifier(tid1)
+    assert(id1 === id2)
+
+    val id3 = HandleIdentifier(id1.publicId, id1.secretId)
+    assert(id3 === id1)
+    val id4 = HandleIdentifier()
+    assert(id4 !== id1)
+  }
+}

kyuubi-common/src/test/scala/org/apache/kyuubi/operation/OperationHandleSuite.scala

@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.operation
+
+import org.apache.hive.service.rpc.thrift.TProtocolVersion
+import org.apache.hive.service.rpc.thrift.TProtocolVersion._
+
+import org.apache.kyuubi.KyuubiFunSuite
+import org.apache.kyuubi.operation.OperationType._
+
+class OperationHandleSuite extends KyuubiFunSuite {
+
+  test("OperationHandle") {
+    val h1 = OperationHandle(EXECUTE_STATEMENT, HIVE_CLI_SERVICE_PROTOCOL_V10)
+    TProtocolVersion.values().foreach { protocol =>
+      assert(h1 === OperationHandle(h1.identifier, h1.typ, protocol))
+    }
+
+    val t1 = h1.toTOperationHandle
+    TProtocolVersion.values().foreach { protocol =>
+      assert(h1 === OperationHandle(t1, protocol))
+    }
+    val h2 = OperationHandle(t1)
+    assert(h1 === h2)
+    assert(!t1.isHasResultSet)
+    h1.setHasResultSet(true)
+    assert(h1.toTOperationHandle.isHasResultSet)
+  }
+}