srielau commented on code in PR #52334:
URL: https://github.com/apache/spark/pull/52334#discussion_r2418049361
##########
sql/api/src/main/scala/org/apache/spark/sql/catalyst/util/SparkParserUtils.scala:
##########
@@ -187,25 +187,93 @@ trait SparkParserUtils {
* Register the origin of the context. Any TreeNode created in the closure
will be assigned the
* registered origin. This method restores the previously set origin after
completion of the
* closure.
+ *
+ * This method is parameter substitution-aware. If parameter substitution
occurred before
+ * parsing, it will automatically adjust the positions and SQL text to refer
to the original SQL
+ * (before substitution) instead of the substituted SQL.
*/
def withOrigin[T](ctx: ParserRuleContext, sqlText: Option[String] = None)(f:
=> T): T = {
val current = CurrentOrigin.get
val text = sqlText.orElse(current.sqlText)
+
if (text.isEmpty) {
CurrentOrigin.set(position(ctx.getStart))
} else {
- CurrentOrigin.set(
- positionAndText(
- ctx.getStart,
- ctx.getStop,
- text.get,
- current.objectType,
- current.objectName))
+ // Check if parameter substitution occurred and adjust origin
accordingly.
+ val adjustedOrigin = adjustOriginForParameterSubstitution(
+ ctx.getStart,
+ ctx.getStop,
+ text.get,
+ current.objectType,
+ current.objectName)
+
+ // Preserve any existing callback when setting the new origin.
+ val finalOrigin = if (current.parameterSubstitutionCallback.isDefined) {
+ adjustedOrigin.copy(parameterSubstitutionCallback =
current.parameterSubstitutionCallback)
+ } else {
+ adjustedOrigin
+ }
+
+ CurrentOrigin.set(finalOrigin)
}
try {
f
} finally {
- CurrentOrigin.set(current)
+ // When restoring origin, preserve any callback that was added during
parsing.
+ val currentAfterParsing = CurrentOrigin.get
+ val originToRestore =
+ if (currentAfterParsing.parameterSubstitutionCallback.isDefined ||
+ current.parameterSubstitutionCallback.isDefined) {
+ // Either the current or the original has a callback - preserve it.
+ val callbackToPreserve =
currentAfterParsing.parameterSubstitutionCallback
+ .orElse(current.parameterSubstitutionCallback)
+ current.copy(parameterSubstitutionCallback = callbackToPreserve)
+ } else {
+ // Neither has a callback - restore as normal.
+ current
+ }
+
+ CurrentOrigin.set(originToRestore)
+ }
+ }
+
+ /**
+ * Adjust origin information to account for parameter substitution.
+ *
+ * If parameter substitution occurred, this method maps positions from the
substituted SQL back
+ * to the original SQL and uses the original SQL text in the origin.
+ *
+ * @param startToken
+ * The start token from the substituted SQL
+ * @param stopToken
+ * The stop token from the substituted SQL
+ * @param substitutedSql
+ * The SQL text after substitution
+ * @param objectType
+ * The object type for the origin
+ * @param objectName
+ * The object name for the origin
+ * @return
+ * Origin with positions and text adjusted for parameter substitution
+ */
+ private def adjustOriginForParameterSubstitution(
+ startToken: Token,
+ stopToken: Token,
+ substitutedSql: String,
+ objectType: Option[String],
+ objectName: Option[String]): Origin = {
+
+ // Try to get parameter substitution callback from CurrentOrigin.
+ CurrentOrigin.get.parameterSubstitutionCallback match {
+ case Some(callback) =>
+ // Cast the callback from Any to the proper type.
+ val typedCallback =
callback.asInstanceOf[SparkParserUtils.ParameterSubstitutionCallback]
Review Comment:
Outdated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
