From 0e68678f2030957a72f4199d8bdc0a37e4dbf2ac Mon Sep 17 00:00:00 2001
From: hongdongdong <hongdongdong@cmss.chinamobile.com>
Date: Thu, 30 Sep 2021 17:53:44 +0800
Subject: [PATCH] [KYUUBI #1184] Support ending with .* to restrict and ingore
 session config

<!--
Thanks for sending a pull request!

Here are some tips for you:
  1. If this is your first time, please read our contributor guidelines: https://kyuubi.readthedocs.io/en/latest/community/contributions.html
  2. If the PR is related to an issue in https://github.com/apache/incubator-kyuubi/issues, add '[KYUUBI #XXXX]' in your PR title, e.g., '[KYUUBI #XXXX] Your PR title ...'.
  3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][KYUUBI #XXXX] Your PR title ...'.
-->

### _Why are the changes needed?_
<!--
Please clarify why the changes are needed. For instance,
  1. If you add a feature, you can talk about the use case of it.
  2. If you fix a bug, you can clarify why it is a bug.
-->
Now, restricted and ingored keys must be complete. We can support eding with '.*' to restricted a list of keys.

### _How was this patch tested?_
- [X] Add some test cases that check the changes thoroughly including negative and positive cases if possible

- [ ] Add screenshots for manual tests if appropriate

- [ ] [Run test](https://kyuubi.readthedocs.io/en/latest/develop_tools/testing.html#running-tests) locally before make a pull request

Closes #1185 from hddong/support-perfix-ignore.

Closes #1184

dade3301 [hongdongdong] reset
c58b0aa6 [hongdongdong] [KYUUBI#1184] Support ending with .* to restrict and ingore session config

Authored-by: hongdongdong <hongdongdong@cmss.chinamobile.com>
Signed-off-by: ulysses-you <ulyssesyou@apache.org>
---
 .../kyuubi/session/SessionManager.scala       | 10 ++++++++--
 .../kyuubi/session/SessionManagerSuite.scala  | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
index 81eee2239..a18576240 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/session/SessionManager.scala
@@ -119,6 +119,10 @@ abstract class SessionManager(name: String) extends CompositeService(name) {
 
   private var _confRestrictList: Set[String] = _
   private var _confIgnoreList: Set[String] = _
+  private lazy val _confRestrictMatchList: Set[String] =
+    _confRestrictList.filter(_.endsWith(".*")).map(_.stripSuffix(".*"))
+  private lazy val _confIgnoreMatchList: Set[String] =
+    _confIgnoreList.filter(_.endsWith(".*")).map(_.stripSuffix(".*"))
 
   // strip prefix and validate whether if key is restricted, ignored or valid
   def validateKey(key: String, value: String): Option[(String, String)] = {
@@ -141,10 +145,12 @@ abstract class SessionManager(name: String) extends CompositeService(name) {
       key
     }
 
-    if (_confRestrictList.contains(normalizedKey)) {
+    if (_confRestrictMatchList.exists(normalizedKey.startsWith(_)) ||
+      _confRestrictList.contains(normalizedKey)) {
       throw KyuubiSQLException(s"$normalizedKey is a restrict key according to the server-side" +
         s" configuration, please remove it and retry if you want to proceed")
-    } else if (_confIgnoreList.contains(normalizedKey)) {
+    } else if (_confIgnoreMatchList.exists(normalizedKey.startsWith(_)) ||
+      _confIgnoreList.contains(normalizedKey)) {
       warn(s"$normalizedKey is a ignored key according to the server-side configuration")
       None
     } else {
diff --git a/kyuubi-common/src/test/scala/org/apache/kyuubi/session/SessionManagerSuite.scala b/kyuubi-common/src/test/scala/org/apache/kyuubi/session/SessionManagerSuite.scala
index 3e8c89301..95432d056 100644
--- a/kyuubi-common/src/test/scala/org/apache/kyuubi/session/SessionManagerSuite.scala
+++ b/kyuubi-common/src/test/scala/org/apache/kyuubi/session/SessionManagerSuite.scala
@@ -23,6 +23,7 @@ import org.apache.hive.service.rpc.thrift._
 import org.scalatest.concurrent.Eventually
 import org.scalatest.time.{Seconds, Span}
 
+import org.apache.kyuubi.KyuubiSQLException
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.service.ThriftFrontendServiceSuite
 
@@ -34,6 +35,8 @@ class SessionManagerSuite extends ThriftFrontendServiceSuite with Eventually {
     .set(KyuubiConf.SESSION_CHECK_INTERVAL, Duration.ofSeconds(5).toMillis)
     .set(KyuubiConf.SESSION_IDLE_TIMEOUT, Duration.ofSeconds(5).toMillis)
     .set(KyuubiConf.OPERATION_IDLE_TIMEOUT, Duration.ofSeconds(20).toMillis)
+    .set(KyuubiConf.SESSION_CONF_RESTRICT_LIST, Seq("spark.*"))
+    .set(KyuubiConf.SESSION_CONF_IGNORE_LIST, Seq("session.engine.*"))
 
   test("close expired operations") {
     withSessionHandle{ (client, handle) =>
@@ -72,4 +75,20 @@ class SessionManagerSuite extends ThriftFrontendServiceSuite with Eventually {
       assert(sessionManager.getOpenSessionCount == 0)
     }
   }
+
+  test("test validate and normalize config") {
+    val sessionManager = server.backendService.sessionManager
+    // test restrict
+    intercept[KyuubiSQLException] {
+      sessionManager.validateAndNormalizeConf(Map("spark.driver.memory" -> "2G"))
+    }
+
+    // test ignore
+    val conf = sessionManager.validateAndNormalizeConf(
+      Map(
+        "session.engine.spark.main.resource" -> "org.apahce.kyuubi.test",
+        "session.check.interval" -> "10000"))
+    assert(conf.size == 1)
+    assert(conf("session.check.interval") == "10000")
+  }
 }