From d31c45a000246f58275ab50ece85b5b2518061a3 Mon Sep 17 00:00:00 2001
From: v-jizhang <jzhang@openteams.com>
Date: Fri, 4 Mar 2022 13:19:58 -0800
Subject: [PATCH] Make sure that quoted identifiers remain quoted in formatter

Cherry-pick of https://github.com/trinodb/trino/pull/11171
The SqlFormatter effectively drops quotes from the quoted identifiers
when formatting 'SET SESSION' statement.
This causes valid queries to fail: e.g. SET SESSION
"name-suffix".property = 'value' PrestoException: Formatted query does
not parse even when identifier (name-suffix) is quoted.

Follow up https://github.com/prestodb/presto/pull/17186

Co-authored-by: Sergey Melnychuk <sergey.melnychuk@starburstdata.com>
---
 .../java/com/facebook/presto/sql/SqlFormatter.java  |  4 ++--
 .../facebook/presto/sql/parser/TestSqlParser.java   | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/SqlFormatter.java b/presto-parser/src/main/java/com/facebook/presto/sql/SqlFormatter.java
index a8adcfc42b780..3be1eaba80ef2 100644
--- a/presto-parser/src/main/java/com/facebook/presto/sql/SqlFormatter.java
+++ b/presto-parser/src/main/java/com/facebook/presto/sql/SqlFormatter.java
@@ -1245,7 +1245,7 @@ protected Void visitInsert(Insert node, Integer indent)
         public Void visitSetSession(SetSession node, Integer context)
         {
             builder.append("SET SESSION ")
-                    .append(node.getName())
+                    .append(formatName(node.getName()))
                     .append(" = ")
                     .append(formatExpression(node.getValue(), parameters));
 
@@ -1256,7 +1256,7 @@ public Void visitSetSession(SetSession node, Integer context)
         public Void visitResetSession(ResetSession node, Integer context)
         {
             builder.append("RESET SESSION ")
-                    .append(node.getName());
+                    .append(formatName(node.getName()));
 
             return null;
         }
diff --git a/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java b/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java
index e97a9655c4334..5e3ec73854a20 100644
--- a/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java
+++ b/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java
@@ -696,6 +696,19 @@ public void testResetSession()
         assertStatement("RESET SESSION foo", new ResetSession(QualifiedName.of("foo")));
     }
 
+    @Test
+    public void testSessionIdentifiers()
+    {
+        assertStatement("SET SESSION \"foo-bar\".baz = 'x'",
+                new SetSession(QualifiedName.of("foo-bar", "baz"), new StringLiteral("x")));
+        assertInvalidStatement("SET SESSION foo-bar.name = 'value'",
+                "mismatched input '-'. Expecting: '.', '='");
+        assertStatement("RESET SESSION \"foo-bar\".baz",
+                new ResetSession(QualifiedName.of("foo-bar", "baz")));
+        assertInvalidStatement("RESET SESSION foo-bar.name",
+                "mismatched input '-'. Expecting: '.', <EOF>");
+    }
+
     @Test
     public void testShowSession()
     {