[presto][verifier] Replace regex-based json_parse safety wrapper with AST-level rewriter (#27202)

Summary:

testing infra fix:
```
com.facebook.presto.verifier.framework.PrestoQueryException: com.facebook.presto.spi.PrestoException: Cannot convert '' to JSON
```
The previous applyJsonParseSafetyWrapper() used SqlFormatter.formatSql() -> regex ->
sqlParser.createStatement() round-trip that silently failed on complex queries with
lambdas and $internal$try(BIND(...)) patterns, leaving bare json_parse() calls that
crash on empty strings during verifier replay.

Replace with AST-level JsonParseTryWrapper using DefaultTreeRewriter +
ExpressionTreeRewriter<Boolean> (same pattern as FunctionCallRewriter). Wraps
json_parse FunctionCall nodes directly in TryExpression on the AST, eliminating the
fragile format/re-parse cycle.

The feature is gated behind a `json-parse-safety-wrapper-enabled` config flag
(default false). Enabled in sapphire velox shadow testing via run-shadow-test.sh.

# Release Notes
```
== NO RELEASE NOTE ==
```

Differential Revision: D94175149

Author: han-yan01

presto-verifier/src/main/java/com/facebook/presto/verifier/framework/VerifierConfig.java

@@ -70,6 +70,8 @@ public class VerifierConfig
     private boolean extendedVerification;
     private String runningMode = CONTROL_TEST_MODE;
 
+    private boolean jsonParseSafetyWrapperEnabled;
+
     private Multimap<String, FunctionCallSubstitute> functionSubstitutes = ImmutableMultimap.of();
 
     @NotNull
@@ -455,6 +457,19 @@ public VerifierConfig setRunningMode(String runningMode)
         return this;
     }
 
+    public boolean isJsonParseSafetyWrapperEnabled()
+    {
+        return jsonParseSafetyWrapperEnabled;
+    }
+
+    @ConfigDescription("Wrap bare json_parse() calls with TRY() during query rewriting to prevent failures on malformed JSON")
+    @Config("json-parse-safety-wrapper-enabled")
+    public VerifierConfig setJsonParseSafetyWrapperEnabled(boolean jsonParseSafetyWrapperEnabled)
+    {
+        this.jsonParseSafetyWrapperEnabled = jsonParseSafetyWrapperEnabled;
+        return this;
+    }
+
     public Multimap<String, FunctionCallSubstitute> getFunctionSubstitutes()
     {
         return functionSubstitutes;

presto-verifier/src/main/java/com/facebook/presto/verifier/rewrite/QueryRewriter.java

@@ -23,7 +23,6 @@
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
-import com.facebook.presto.sql.SqlFormatter;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.planner.LiteralEncoder;
 import com.facebook.presto.sql.tree.AllColumns;
@@ -36,12 +35,15 @@
 import com.facebook.presto.sql.tree.DropTable;
 import com.facebook.presto.sql.tree.DropView;
 import com.facebook.presto.sql.tree.Expression;
+import com.facebook.presto.sql.tree.ExpressionRewriter;
+import com.facebook.presto.sql.tree.ExpressionTreeRewriter;
 import com.facebook.presto.sql.tree.FunctionCall;
 import com.facebook.presto.sql.tree.Identifier;
 import com.facebook.presto.sql.tree.Insert;
 import com.facebook.presto.sql.tree.IsNullPredicate;
 import com.facebook.presto.sql.tree.LikeClause;
 import com.facebook.presto.sql.tree.LogicalBinaryExpression;
+import com.facebook.presto.sql.tree.Node;
 import com.facebook.presto.sql.tree.Property;
 import com.facebook.presto.sql.tree.QualifiedName;
 import com.facebook.presto.sql.tree.Query;
@@ -51,9 +53,10 @@
 import com.facebook.presto.sql.tree.ShowCreate;
 import com.facebook.presto.sql.tree.SingleColumn;
 import com.facebook.presto.sql.tree.Statement;
+import com.facebook.presto.sql.tree.SubqueryExpression;
+import com.facebook.presto.sql.tree.TryExpression;
 import com.facebook.presto.verifier.framework.ClusterType;
 import com.facebook.presto.verifier.framework.Column;
-import com.facebook.presto.verifier.framework.JsonParseSafetyWrapper;
 import com.facebook.presto.verifier.framework.QueryConfiguration;
 import com.facebook.presto.verifier.framework.QueryException;
 import com.facebook.presto.verifier.framework.QueryObjectBundle;
@@ -111,8 +114,6 @@
 
 public class QueryRewriter
 {
-    private static final com.facebook.airlift.log.Logger log = com.facebook.airlift.log.Logger.get(QueryRewriter.class);
-
     private final SqlParser sqlParser;
     private final TypeManager typeManager;
     private final BlockEncodingSerde blockEncodingSerde;
@@ -121,6 +122,7 @@ public class QueryRewriter
     private final Map<ClusterType, List<Property>> tableProperties;
     private final Map<ClusterType, Boolean> reuseTables;
     private final Optional<FunctionCallRewriter> functionCallRewriter;
+    private final boolean jsonParseSafetyWrapperEnabled;
 
     public QueryRewriter(
             SqlParser sqlParser,
@@ -131,7 +133,7 @@ public QueryRewriter(
             Map<ClusterType, List<Property>> tableProperties,
             Map<ClusterType, Boolean> reuseTables)
     {
-        this(sqlParser, typeManager, blockEncodingSerde, prestoAction, tablePrefixes, tableProperties, reuseTables, ImmutableMultimap.of());
+        this(sqlParser, typeManager, blockEncodingSerde, prestoAction, tablePrefixes, tableProperties, reuseTables, ImmutableMultimap.of(), false);
     }
 
     public QueryRewriter(
@@ -143,6 +145,20 @@ public QueryRewriter(
             Map<ClusterType, List<Property>> tableProperties,
             Map<ClusterType, Boolean> reuseTables,
             Multimap<String, FunctionCallSubstitute> functionSubstitutes)
+    {
+        this(sqlParser, typeManager, blockEncodingSerde, prestoAction, tablePrefixes, tableProperties, reuseTables, functionSubstitutes, false);
+    }
+
+    public QueryRewriter(
+            SqlParser sqlParser,
+            TypeManager typeManager,
+            BlockEncodingSerde blockEncodingSerde,
+            PrestoAction prestoAction,
+            Map<ClusterType, QualifiedName> tablePrefixes,
+            Map<ClusterType, List<Property>> tableProperties,
+            Map<ClusterType, Boolean> reuseTables,
+            Multimap<String, FunctionCallSubstitute> functionSubstitutes,
+            boolean jsonParseSafetyWrapperEnabled)
     {
         this.sqlParser = requireNonNull(sqlParser, "sqlParser is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
@@ -152,26 +168,64 @@ public QueryRewriter(
         this.tableProperties = ImmutableMap.copyOf(tableProperties);
         this.reuseTables = ImmutableMap.copyOf(reuseTables);
         this.functionCallRewriter = FunctionCallRewriter.getInstance(functionSubstitutes, typeManager);
+        this.jsonParseSafetyWrapperEnabled = jsonParseSafetyWrapperEnabled;
     }
 
     /**
-     * Helper method to apply json_parse safety wrapper to a query.
-     * Wraps json_parse() calls with TRY() to handle malformed JSON gracefully.
-     * If re-parsing fails, returns the original query unchanged.
+     * Wraps bare json_parse() calls with TRY() using AST rewriting.
+     * Operates directly on the AST to avoid fragile SQL format/re-parse round-trips
+     * that silently fail on complex queries with lambdas or internal function patterns.
      */
-    private Query applyJsonParseSafetyWrapper(Query query)
+    private static Query applyJsonParseSafetyWrapper(Query query)
     {
-        try {
-            String sql = SqlFormatter.formatSql(query, Optional.empty());
-            String fixedSql = JsonParseSafetyWrapper.wrapUnsafeJsonParse(sql);
-            if (!sql.equals(fixedSql)) {
-                return (Query) sqlParser.createStatement(fixedSql, PARSING_OPTIONS);
-            }
-        }
-        catch (Exception e) {
-            log.warn(e, "Failed to apply json_parse safety wrapper, using original query");
+        return (Query) new JsonParseTryWrapper().process(query, null);
+    }
+
+    /**
+     * AST-level rewriter that wraps json_parse() FunctionCall nodes in TryExpression.
+     * Uses the same DefaultTreeRewriter + ExpressionTreeRewriter pattern as FunctionCallRewriter.
+     */
+    private static class JsonParseTryWrapper
+            extends DefaultTreeRewriter<Void>
+    {
+        @Override
+        protected Node visitExpression(Expression node, Void context)
+        {
+            JsonParseTryWrapper statementRewriter = this;
+
+            return ExpressionTreeRewriter.rewriteWith(new ExpressionRewriter<Boolean>()
+            {
+                @Override
+                public Expression rewriteFunctionCall(FunctionCall original, Boolean insideTry, ExpressionTreeRewriter<Boolean> treeRewriter)
+                {
+                    FunctionCall defaultRewrite = treeRewriter.defaultRewrite(original, false);
+                    if (defaultRewrite.getName().getSuffix().equalsIgnoreCase("json_parse") && !Boolean.TRUE.equals(insideTry)) {
+                        return new TryExpression(defaultRewrite);
+                    }
+                    return defaultRewrite;
+                }
+
+                @Override
+                public Expression rewriteTryExpression(TryExpression original, Boolean insideTry, ExpressionTreeRewriter<Boolean> treeRewriter)
+                {
+                    Expression inner = treeRewriter.rewrite(original.getInnerExpression(), true);
+                    if (inner != original.getInnerExpression()) {
+                        return new TryExpression(inner);
+                    }
+                    return original;
+                }
+
+                @Override
+                public Expression rewriteSubqueryExpression(SubqueryExpression expression, Boolean insideTry, ExpressionTreeRewriter<Boolean> treeRewriter)
+                {
+                    Node rewritten = statementRewriter.process(expression.getQuery(), null);
+                    if (expression.getQuery() == rewritten) {
+                        return expression;
+                    }
+                    return new SubqueryExpression((Query) rewritten);
+                }
+            }, node, false);
         }
-        return query;
     }
 
     public QueryObjectBundle rewriteQuery(@Language("SQL") String query, QueryConfiguration queryConfiguration, ClusterType clusterType)
@@ -197,8 +251,9 @@ public QueryObjectBundle rewriteQuery(@Language("SQL") String query, QueryConfig
                 FunctionCallRewriter.RewriterResult rewriterResult = functionCallRewriter.get().rewrite(createQuery);
                 createQuery = (Query) rewriterResult.getRewrittenNode();
                 functionSubstitutions = rewriterResult.getSubstitutions();
-
-                // Apply safety wrapper for json_parse() calls to prevent failures on malformed JSON
+            }
+            // Apply safety wrapper for json_parse() calls to prevent failures on malformed JSON
+            if (jsonParseSafetyWrapperEnabled) {
                 createQuery = applyJsonParseSafetyWrapper(createQuery);
             }
             if (shouldReuseTable && !functionSubstitutions.isPresent()) {
@@ -244,8 +299,9 @@ public QueryObjectBundle rewriteQuery(@Language("SQL") String query, QueryConfig
                 FunctionCallRewriter.RewriterResult rewriterResult = functionCallRewriter.get().rewrite(insertQuery);
                 insertQuery = (Query) rewriterResult.getRewrittenNode();
                 functionSubstitutions = rewriterResult.getSubstitutions();
-
-                // Apply safety wrapper for json_parse() calls
+            }
+            // Apply safety wrapper for json_parse() calls
+            if (jsonParseSafetyWrapperEnabled) {
                 insertQuery = applyJsonParseSafetyWrapper(insertQuery);
             }
             if (shouldReuseTable && !functionSubstitutions.isPresent()) {
@@ -291,8 +347,9 @@ public QueryObjectBundle rewriteQuery(@Language("SQL") String query, QueryConfig
                 FunctionCallRewriter.RewriterResult rewriterResult = functionCallRewriter.get().rewrite(queryBody);
                 queryBody = (Query) rewriterResult.getRewrittenNode();
                 functionSubstitutions = rewriterResult.getSubstitutions();
-
-                // Apply safety wrapper for json_parse() calls
+            }
+            // Apply safety wrapper for json_parse() calls
+            if (jsonParseSafetyWrapperEnabled) {
                 queryBody = applyJsonParseSafetyWrapper(queryBody);
             }
 

presto-verifier/src/main/java/com/facebook/presto/verifier/rewrite/VerificationQueryRewriterFactory.java

@@ -57,6 +57,7 @@ public class VerificationQueryRewriterFactory
     private final boolean testReuseTable;
 
     private final Multimap<String, FunctionCallSubstitute> functionSubstitutes;
+    private final boolean jsonParseSafetyWrapperEnabled;
 
     @Inject
     public VerificationQueryRewriterFactory(
@@ -77,6 +78,7 @@ public VerificationQueryRewriterFactory(
         this.controlReuseTable = controlConfig.isReuseTable();
         this.testReuseTable = testConfig.isReuseTable();
         this.functionSubstitutes = verifierConfig.getFunctionSubstitutes();
+        this.jsonParseSafetyWrapperEnabled = verifierConfig.isJsonParseSafetyWrapperEnabled();
     }
 
     @Override
@@ -90,7 +92,8 @@ public QueryRewriter create(PrestoAction prestoAction)
                 ImmutableMap.of(CONTROL, controlTablePrefix, TEST, testTablePrefix),
                 ImmutableMap.of(CONTROL, controlTableProperties, TEST, testTableProperties),
                 ImmutableMap.of(CONTROL, controlReuseTable, TEST, testReuseTable),
-                functionSubstitutes);
+                functionSubstitutes,
+                jsonParseSafetyWrapperEnabled);
     }
 
     private static List<Property> constructProperties(Map<String, Object> propertiesMap)

presto-verifier/src/test/java/com/facebook/presto/verifier/rewrite/TestQueryRewriter.java

@@ -84,7 +84,8 @@ public class TestQueryRewriter
     private static final QueryRewriteConfig QUERY_REWRITE_CONFIG = new QueryRewriteConfig()
             .setTablePrefix("local.tmp")
             .setTableProperties("{\"p_int\": 30, \"p_long\": 4294967297, \"p_double\": 1.5, \"p_varchar\": \"test\", \"p_bool\": true}");
-    private static final VerifierConfig VERIFIER_CONFIG = new VerifierConfig();
+    private static final VerifierConfig VERIFIER_CONFIG = new VerifierConfig()
+            .setJsonParseSafetyWrapperEnabled(true);
     private static final SqlParser SQL_PARSER = new SqlParser(new SqlParserOptions().allowIdentifierSymbol(COLON, AT_SIGN));
     private static final BlockEncodingSerde blockEncodingSerde = new BlockEncodingManager();
 
@@ -668,6 +669,106 @@ private static void assertStatements(List<Statement> actual, List<Statement> exp
         assertEquals(actualQueries, expectedQueries);
     }
 
+    @Test
+    public void testJsonParseSafetyWrapper()
+    {
+        QueryRewriter queryRewriter = getQueryRewriter();
+
+        // Bare json_parse should be wrapped in TRY
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT json_parse(b) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(b)) FROM test_table");
+
+        // Already wrapped in TRY should not be double-wrapped
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT TRY(json_parse(b)) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(b)) FROM test_table");
+
+        // Mixed: some wrapped, some not
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT TRY(json_parse(b)), json_parse(b) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(b)), TRY(json_parse(b)) FROM test_table");
+
+        // json_parse inside complex expression with lambda
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT transform(ARRAY[b], x -> json_parse(x)) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT transform(ARRAY[b], x -> TRY(json_parse(x))) FROM test_table");
+
+        // json_parse in subquery
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT * FROM (SELECT json_parse(b) AS parsed FROM test_table) t",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT * FROM (SELECT TRY(json_parse(b)) AS parsed FROM test_table) t");
+
+        // json_parse in WHERE clause
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT a FROM test_table WHERE json_parse(b) IS NOT NULL",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT a FROM test_table WHERE TRY(json_parse(b)) IS NOT NULL");
+
+        // json_parse nested inside json_extract (common pattern from function substitution)
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT json_extract(json_parse(CAST(b AS varchar)), '$.key') FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT json_extract(TRY(json_parse(CAST(b AS varchar))), '$.key') FROM test_table");
+
+        // No json_parse - should pass through unchanged
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT a, b FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT a, b FROM test_table");
+
+        // Multiple json_parse in SELECT + WHERE (both should be wrapped independently)
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT json_parse(b) FROM test_table WHERE CAST(json_parse(b) AS VARCHAR) <> ''",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(b)) FROM test_table WHERE CAST(TRY(json_parse(b)) AS VARCHAR) <> ''");
+
+        // COALESCE(TRY(json_parse(col)), JSON 'null') - json_parse already inside TRY, no wrapping
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT COALESCE(TRY(json_parse(b)), JSON 'null') FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT COALESCE(TRY(json_parse(b)), JSON 'null') FROM test_table");
+
+        // Many json_parse calls in a single query (stress test for AST approach vs regex)
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT json_parse(b), TRY(json_parse(b)), json_parse(CONCAT(b, '')), json_parse(UPPER(b)), json_parse(LOWER(b)) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(b)), TRY(json_parse(b)), TRY(json_parse(CONCAT(b, ''))), TRY(json_parse(UPPER(b))), TRY(json_parse(LOWER(b))) FROM test_table");
+    }
+
+    @Test
+    public void testJsonParseSafetyWrapperWithFunctionSubstitutes()
+    {
+        // Test that json_parse wrapping works AFTER function substitution
+        VerifierConfig verifierConfig = new VerifierConfig()
+                .setJsonParseSafetyWrapperEnabled(true)
+                .setFunctionSubstitutes("/arbitrary(x)/min(x)/");
+        QueryRewriter queryRewriter = getQueryRewriter(new QueryRewriteConfig(), verifierConfig);
+
+        // Function substitution + json_parse wrapping should both apply
+        assertCreateTableAs(
+                queryRewriter.rewriteQuery(
+                        "SELECT json_parse(ARBITRARY(b)) FROM test_table",
+                        CONFIGURATION, CONTROL).getQuery(),
+                "SELECT TRY(json_parse(MIN(b))) FROM test_table");
+    }
+
     private QueryRewriter getQueryRewriter()
     {
         return getQueryRewriter(QUERY_REWRITE_CONFIG, VERIFIER_CONFIG);