diff --git a/src/main/scala/org/apache/spark/KyuubiSparkUtil.scala b/src/main/scala/org/apache/spark/KyuubiSparkUtil.scala index c6a76bdfc..861f879ed 100644 --- a/src/main/scala/org/apache/spark/KyuubiSparkUtil.scala +++ b/src/main/scala/org/apache/spark/KyuubiSparkUtil.scala @@ -123,4 +123,12 @@ object KyuubiSparkUtil extends Logging { val sMinor = minorVersion(version) tMajor > sMajor || (tMajor == sMajor && tMinor >= sMinor) } + + /** + * Convert a time parameter such as (50s, 100ms, or 250us) to microseconds for internal use. If + * no suffix is provided, the passed number is assumed to be in ms. + */ + def timeStringAsMs(str: String): Long = { + Utils.timeStringAsMs(str) + } } diff --git a/src/main/scala/yaooqinn/kyuubi/operation/KyuubiOperation.scala b/src/main/scala/yaooqinn/kyuubi/operation/KyuubiOperation.scala index 0f3a37692..0efe1e0a7 100644 --- a/src/main/scala/yaooqinn/kyuubi/operation/KyuubiOperation.scala +++ b/src/main/scala/yaooqinn/kyuubi/operation/KyuubiOperation.scala @@ -49,7 +49,8 @@ class KyuubiOperation(session: KyuubiSession, statement: String) extends Logging private[this] val conf = session.sparkSession.conf - private[this] val operationTimeout = conf.get(OPERATION_IDLE_TIMEOUT.key).toLong + private[this] val operationTimeout = + KyuubiSparkUtil.timeStringAsMs(conf.get(OPERATION_IDLE_TIMEOUT.key)) private[this] var lastAccessTime = System.currentTimeMillis() private[this] var hasResultSet: Boolean = false diff --git a/src/test/scala/org/apache/spark/KyuubiSparkUtilSuite.scala b/src/test/scala/org/apache/spark/KyuubiSparkUtilSuite.scala index f5372414f..e8415b6aa 100644 --- a/src/test/scala/org/apache/spark/KyuubiSparkUtilSuite.scala +++ b/src/test/scala/org/apache/spark/KyuubiSparkUtilSuite.scala @@ -148,4 +148,10 @@ class KyuubiSparkUtilSuite extends SparkFunSuite { assert(!KyuubiSparkUtil.equalOrHigherThan("2.4.1")) assert(!KyuubiSparkUtil.equalOrHigherThan("3.0.0")) } + + test("testTimeStringAsMs") { + assert(KyuubiSparkUtil.timeStringAsMs("50s") === 50000L) + assert(KyuubiSparkUtil.timeStringAsMs("50min") === 50 * 60 * 1000L) + assert(KyuubiSparkUtil.timeStringAsMs("100ms") === 100L) + } }