SUBMARINE-638. [Spark-Security]Apply global/local limit after masking projection

submarine-security/spark-security/src/main/scala/org/apache/spark/sql/catalyst/optimizer/SubmarineDataMaskingExtension.scala

@@ -28,7 +28,7 @@ import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.{FunctionIdentifier, TableIdentifier}
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.catalog.{CatalogFunction, CatalogTable, HiveTableRelation}
-import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, ExprId, NamedExpression, SubqueryExpression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression, ExprId, NamedExpression, SubqueryExpression}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.execution.command.{CreateDataSourceTableAsSelectCommand, CreateViewCommand, InsertIntoDataSourceDirCommand}
@@ -233,7 +233,16 @@ case class SubmarineDataMaskingExtension(spark: SparkSession) extends Rule[Logic
         case p if hasCatalogTable(p) => SubmarineDataMasking(p)
       }
 
-      marked transformAllExpressions {
+      // Extract global/local limit if any and apply after masking projection
+      val limitExpr: Option[Expression] = plan match {
+        case globalLimit: GlobalLimit => Some(globalLimit.limitExpr)
+        case localLimit: LocalLimit => Some(localLimit.limitExpr)
+        case _ => None
+      }
+
+      val markedWithLimit = if (limitExpr.isDefined) Limit(limitExpr.get, marked) else marked
+
+      markedWithLimit transformAllExpressions {
         case s: SubqueryExpression =>
           val SubqueryCompatible(newPlan, _) = SubqueryCompatible(
               SubmarineDataMasking(s.plan), SubqueryExpression.hasCorrelatedSubquery(s))

submarine-security/spark-security/src/test/scala/org/apache/submarine/spark/security/DataMaskingSQLTest.scala

@@ -21,7 +21,7 @@ package org.apache.submarine.spark.security
 
 import org.apache.commons.codec.digest.DigestUtils
 import org.apache.spark.sql.SubmarineSparkUtils.{enableDataMasking, withUser}
-import org.apache.spark.sql.catalyst.plans.logical.{Project, SubmarineDataMasking}
+import org.apache.spark.sql.catalyst.plans.logical.{GlobalLimit, Project, SubmarineDataMasking}
 import org.apache.spark.sql.hive.test.TestHive
 import org.scalatest.{BeforeAndAfterAll, FunSuite}
 
@@ -282,4 +282,14 @@ case class DataMaskingSQLTest() extends FunSuite with BeforeAndAfterAll {
       assert(row.getString(1) === "xxx_277", "value shows last 4 characters")
     }
   }
+
+  test("query limit expression") {
+    withUser("bob") {
+      val df = sql("select * from default.src limit 10")
+      assert(df.queryExecution.optimizedPlan.find(_.isInstanceOf[SubmarineDataMasking]).nonEmpty)
+      assert(df.queryExecution.optimizedPlan.isInstanceOf[GlobalLimit])
+      val row = df.take(1)(0)
+      assert(row.getString(1).startsWith("x"), "values should be masked")
+    }
+  }
 }