[SPARK-48530][SQL] Support for local variables in SQL Scripting

common/utils/src/main/resources/error/error-conditions.json

@@ -3592,6 +3592,16 @@
         "message" : [
           "Variable <varName> can only be declared at the beginning of the compound."
         ]
+      },
+      "QUALIFIED_LOCAL_VARIABLE" : {
+        "message" : [
+          "The variable <varName> must be declared without a qualifier, as qualifiers are not allowed for local variable declarations."
+        ]
+      },
+      "REPLACE_LOCAL_VARIABLE" : {
+        "message" : [
+          "The variable <varName> does not support DECLARE OR REPLACE, as local variables cannot be replaced."
+        ]
       }
     },
     "sqlState" : "42K0M"
@@ -3738,6 +3748,12 @@
     ],
     "sqlState" : "42K0L"
   },
+  "LABEL_NAME_FORBIDDEN" : {
+    "message" : [
+      "The label name <label> is forbidden."
+    ],
+    "sqlState" : "42K0L"
+  },
   "LOAD_DATA_PATH_NOT_EXISTS" : {
     "message" : [
       "LOAD DATA input path does not exist: <path>."

common/utils/src/main/scala/org/apache/spark/util/LexicalThreadLocal.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.spark.util
+
+/**
+ * Helper trait for defining thread locals with lexical scoping. With this helper, the thread local
+ * is private and can only be set by the [[Handle]]. The [[Handle]] only exposes the thread local
+ * value to functions passed into its [[runWith]] method. This pattern allows for the lifetime of
+ * the thread local value to be strictly controlled.
+ *
+ * Rather than calling `tl.set(...)` and `tl.remove()` you would get a handle and execute your code
+ * in `handle.runWith { ... }`.
+ *
+ * Example:
+ * {{{
+ *   object Credentials extends LexicalThreadLocal[Int] {
+ *     def create(creds: Map[String, String]) = new Handle(Some(creds))
+ *   }
+ *   ...
+ *   val handle = Credentials.create(Map("key" -> "value"))
+ *   assert(Credentials.get() == None)
+ *   handle.runWith {
+ *     assert(Credentials.get() == Some(Map("key" -> "value")))
+ *   }
+ * }}}
+ */
+trait LexicalThreadLocal[T] {
+  private val tl = new ThreadLocal[T]
+
+  private def set(opt: Option[T]): Unit = {
+    opt match {
+      case Some(x) => tl.set(x)
+      case None => tl.remove()
+    }
+  }
+
+  protected def createHandle(opt: Option[T]): Handle = new Handle(opt)
+
+  def get(): Option[T] = Option(tl.get)
+
+  /** Final class representing a handle to a thread local value. */
+  final class Handle private[LexicalThreadLocal] (private val opt: Option[T]) {
+    def runWith[R](f: => R): R = {
+      val old = get()
+      set(opt)
+      try f finally {
+        set(old)
+      }
+    }
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlScriptingLocalVariableManager.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.spark.sql.catalyst
+
+import org.apache.spark.sql.catalyst.catalog.VariableManager
+import org.apache.spark.util.LexicalThreadLocal
+
+object SqlScriptingLocalVariableManager extends LexicalThreadLocal[VariableManager] {
+  def create(variableManager: VariableManager): Handle = createHandle(Option(variableManager))
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -139,6 +139,9 @@ object FakeV2SessionCatalog extends TableCatalog with FunctionCatalog with Suppo
  *                              even if a temp view `t` has been created.
  * @param outerPlan The query plan from the outer query that can be used to resolve star
  *                  expressions in a subquery.
+ * @param isExecuteImmediate Whether the current plan is created by EXECUTE IMMEDIATE. Used when
+ *                           resolving variables, as SQL Scripting local variables should not be
+ *                           visible from EXECUTE IMMEDIATE.
  */
 case class AnalysisContext(
     catalogAndNamespace: Seq[String] = Nil,
@@ -154,6 +157,7 @@ case class AnalysisContext(
     referredTempFunctionNames: mutable.Set[String] = mutable.Set.empty,
     referredTempVariableNames: Seq[Seq[String]] = Seq.empty,
     outerPlan: Option[LogicalPlan] = None,
+    isExecuteImmediate: Boolean = false,
 
     /**
      * This is a bridge state between this fixed-point [[Analyzer]] and a single-pass [[Resolver]].
@@ -208,7 +212,16 @@ object AnalysisContext {
       originContext.relationCache,
       viewDesc.viewReferredTempViewNames,
       mutable.Set(viewDesc.viewReferredTempFunctionNames: _*),
-      viewDesc.viewReferredTempVariableNames)
+      viewDesc.viewReferredTempVariableNames,
+      isExecuteImmediate = originContext.isExecuteImmediate)
+    set(context)
+    try f finally { set(originContext) }
+  }
+
+  def withExecuteImmediateContext[A](f: => A): A = {
+    val originContext = value.get()
+    val context = originContext.copy(isExecuteImmediate = true)
+
     set(context)
     try f finally { set(originContext) }
   }
@@ -325,7 +338,10 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
 
   override def batches: Seq[Batch] = Seq(
     Batch("Substitution", fixedPoint,
-      new SubstituteExecuteImmediate(catalogManager),
+      new SubstituteExecuteImmediate(
+        catalogManager,
+        resolveChild = executeSameContext,
+        checkAnalysis = checkAnalysis),
       // This rule optimizes `UpdateFields` expression chains so looks more like optimization rule.
       // However, when manipulating deeply nested schema, `UpdateFields` expression tree could be
       // very complex and make analysis impossible. Thus we need to optimize `UpdateFields` early

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ColumnResolutionHelper.scala

@@ -23,6 +23,7 @@ import scala.collection.mutable
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.SqlScriptingLocalVariableManager
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.SubExprUtils.wrapOuterReference
 import org.apache.spark.sql.catalyst.plans.logical._
@@ -251,6 +252,14 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
     }
   }
 
+  /**
+   * Look up variable by nameParts.
+   * If in SQL Script, first check local variables, unless in EXECUTE IMMEDIATE
+   * (EXECUTE IMMEDIATE generated query cannot access local variables).
+   * if not found fall back to session variables.
+   * @param nameParts NameParts of the variable.
+   * @return Reference to the variable.
+   */
   def lookupVariable(nameParts: Seq[String]): Option[VariableReference] = {
     // The temp variables live in `SYSTEM.SESSION`, and the name can be qualified or not.
     def maybeTempVariableName(nameParts: Seq[String]): Boolean = {
@@ -266,22 +275,48 @@ trait ColumnResolutionHelper extends Logging with DataTypeErrorsBase {
       }
     }
 
-    if (maybeTempVariableName(nameParts)) {
-      val variableName = if (conf.caseSensitiveAnalysis) {
-        nameParts.last
-      } else {
-        nameParts.last.toLowerCase(Locale.ROOT)
-      }
-      catalogManager.tempVariableManager.get(variableName).map { varDef =>
+    val namePartsCaseAdjusted = if (conf.caseSensitiveAnalysis) {
+      nameParts
+    } else {
+      nameParts.map(_.toLowerCase(Locale.ROOT))
+    }
+
+    SqlScriptingLocalVariableManager.get()
+      // If sessionOnly is set to true lookup only session variables.
+      .filterNot(_ => AnalysisContext.get.isExecuteImmediate)
+      // If variable name is qualified with system.session.<varName> treat it as a session variable.
+      .filterNot(_ =>
+        nameParts.length == 3
+          && nameParts.take(2).map(_.toLowerCase(Locale.ROOT)) == Seq("system", "session"))
+      // If variable name is qualified with session.<varName> treat it as a session variable.
+      .filterNot(_ =>
+        nameParts.length == 2
+          && nameParts.head.toLowerCase(Locale.ROOT) == "session")
+      // Local variable must be in format <varName> or <label>.<varName>
+      .filter(_ => namePartsCaseAdjusted.nonEmpty && namePartsCaseAdjusted.length <= 2)
+      .flatMap(_.get(namePartsCaseAdjusted))
+      .map { varDef =>
         VariableReference(
           nameParts,
           FakeSystemCatalog,
-          Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), variableName),
+          Identifier.of(Array(varDef.identifier.namespace().last), namePartsCaseAdjusted.last),
           varDef)
       }
-    } else {
-      None
-    }
+      .orElse(
+        if (maybeTempVariableName(nameParts)) {
+          catalogManager.tempVariableManager
+            .get(namePartsCaseAdjusted)
+            .map { varDef =>
+              VariableReference(
+                nameParts,
+                FakeSystemCatalog,
+                Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), namePartsCaseAdjusted.last),
+                varDef
+              )}
+        } else {
+          None
+        }
+      )
   }
 
   // Resolves `UnresolvedAttribute` to its value.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCatalogs.scala

@@ -19,9 +19,12 @@ package org.apache.spark.sql.catalyst.analysis
 
 import scala.jdk.CollectionConverters._
 
+import org.apache.spark.sql.AnalysisException
+import org.apache.spark.sql.catalyst.SqlScriptingLocalVariableManager
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, Identifier, LookupCatalog, SupportsNamespaces}
+import org.apache.spark.sql.errors.DataTypeErrors.toSQLId
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.util.ArrayImplicits._
 
@@ -34,11 +37,30 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperatorsDown {
     // We only support temp variables for now and the system catalog is not properly implemented
     // yet. We need to resolve `UnresolvedIdentifier` for variable commands specially.
-    case c @ CreateVariable(UnresolvedIdentifier(nameParts, _), _, _) =>
-      val resolved = resolveVariableName(nameParts)
-      c.copy(name = resolved)
+    case c @ CreateVariable(UnresolvedIdentifier(nameParts, _), _, _, _) =>
+      // From scripts we can only create local variables, which must be unqualified,
+      // and must not be DECLARE OR REPLACE.
+      if (SqlScriptingLocalVariableManager.get().isDefined &&
+        !AnalysisContext.get.isExecuteImmediate) {
+        // TODO [SPARK-50785]: Uncomment this when For Statement starts properly using local vars.
+//        if (c.replace) {
+//          throw new AnalysisException(
+//            "INVALID_VARIABLE_DECLARATION.REPLACE_LOCAL_VARIABLE",
+//            Map("varName" -> toSQLId(nameParts))
+//          )
+//        }
+
+        if (nameParts.length != 1) {
+          throw new AnalysisException(
+            "INVALID_VARIABLE_DECLARATION.QUALIFIED_LOCAL_VARIABLE",
+            Map("varName" -> toSQLId(nameParts)))
+        }
+      }
+
+      val resolved = resolveCreateVariableName(nameParts)
+      c.copy(name = resolved, sessionVariablesOnly = AnalysisContext.get.isExecuteImmediate)
     case d @ DropVariable(UnresolvedIdentifier(nameParts, _), _) =>
-      val resolved = resolveVariableName(nameParts)
+      val resolved = resolveDropVariableName(nameParts)
       d.copy(name = resolved)
 
     case UnresolvedIdentifier(nameParts, allowTemp) =>
@@ -73,28 +95,40 @@ class ResolveCatalogs(val catalogManager: CatalogManager)
     }
   }
 
-  private def resolveVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
-    def ident: Identifier = Identifier.of(Array(CatalogManager.SESSION_NAMESPACE), nameParts.last)
-    if (nameParts.length == 1) {
+  private def resolveCreateVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
+    val ident = SqlScriptingLocalVariableManager.get()
+      .filterNot(_ => AnalysisContext.get.isExecuteImmediate)
+      .getOrElse(catalogManager.tempVariableManager)
+      .createIdentifier(nameParts.last)
+
+    resolveVariableName(nameParts, ident)
+  }
+
+  private def resolveDropVariableName(nameParts: Seq[String]): ResolvedIdentifier = {
+    // Only session variables can be dropped, so catalogManager.scriptingLocalVariableManager
+    // is not checked in the case of DropVariable.
+    val ident = catalogManager.tempVariableManager.createIdentifier(nameParts.last)
+    resolveVariableName(nameParts, ident)
+  }
+
+  private def resolveVariableName(
+      nameParts: Seq[String],
+      ident: Identifier): ResolvedIdentifier = nameParts.length match {
+    case 1 => ResolvedIdentifier(FakeSystemCatalog, ident)
+
+    // On declare variable, local variables support only unqualified names.
+    // On drop variable, local variables are not supported at all.
+    case 2 if nameParts.head.equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE) =>
       ResolvedIdentifier(FakeSystemCatalog, ident)
-    } else if (nameParts.length == 2) {
-      if (nameParts.head.equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE)) {
-        ResolvedIdentifier(FakeSystemCatalog, ident)
-      } else {
-        throw QueryCompilationErrors.unresolvedVariableError(
-          nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-      }
-    } else if (nameParts.length == 3) {
-      if (nameParts(0).equalsIgnoreCase(CatalogManager.SYSTEM_CATALOG_NAME) &&
-        nameParts(1).equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE)) {
-        ResolvedIdentifier(FakeSystemCatalog, ident)
-      } else {
-        throw QueryCompilationErrors.unresolvedVariableError(
-          nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-      }
-    } else {
+
+    // When there are 3 nameParts the variable must be a fully qualified session variable
+    // i.e. "system.session.<varName>"
+    case 3 if nameParts(0).equalsIgnoreCase(CatalogManager.SYSTEM_CATALOG_NAME) &&
+      nameParts(1).equalsIgnoreCase(CatalogManager.SESSION_NAMESPACE) =>
+      ResolvedIdentifier(FakeSystemCatalog, ident)
+
+    case _ =>
       throw QueryCompilationErrors.unresolvedVariableError(
-        nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, CatalogManager.SESSION_NAMESPACE))
-    }
+        nameParts, Seq(CatalogManager.SYSTEM_CATALOG_NAME, ident.namespace().head))
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSetVariable.scala

@@ -53,14 +53,17 @@ class ResolveSetVariable(val catalogManager: CatalogManager) extends Rule[Logica
       // Names are normalized when the variables are created.
       // No need for case insensitive comparison here.
       // TODO: we need to group by the qualified variable name once other catalogs support it.
-      val dups = resolvedVars.groupBy(_.identifier.name).filter(kv => kv._2.length > 1)
+      val dups = resolvedVars.groupBy(_.identifier).filter(kv => kv._2.length > 1)
       if (dups.nonEmpty) {
         throw new AnalysisException(
           errorClass = "DUPLICATE_ASSIGNMENTS",
-          messageParameters = Map("nameList" -> dups.keys.map(toSQLId).mkString(", ")))
+          messageParameters = Map("nameList" ->
+            dups.keys.map(key => toSQLId(key.name())).mkString(", ")))
       }
 
-      setVariable.copy(targetVariables = resolvedVars)
+      setVariable.copy(
+        targetVariables = resolvedVars,
+        sessionVariablesOnly = AnalysisContext.get.isExecuteImmediate)
 
     case setVariable: SetVariable
         if setVariable.targetVariables.forall(_.isInstanceOf[VariableReference]) &&