Support formatting Use statement

prestodb · Oct 12, 2020 · ac60c75 · ac60c75
1 parent 93149ae
commit ac60c75
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
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
@@ -107,6 +107,7 @@
 import com.facebook.presto.sql.tree.TransactionMode;
 import com.facebook.presto.sql.tree.Union;
 import com.facebook.presto.sql.tree.Unnest;
+import com.facebook.presto.sql.tree.Use;
 import com.facebook.presto.sql.tree.Values;
 import com.facebook.presto.sql.tree.With;
 import com.facebook.presto.sql.tree.WithQuery;
@@ -785,6 +786,19 @@ protected Void visitShowFunctions(ShowFunctions node, Integer context)
             return null;
         }
 
+        @Override
+        protected Void visitUse(Use node, Integer context)
+        {
+            builder.append("USE ");
+            if (node.getCatalog().isPresent()) {
+                builder.append(formatExpression(node.getCatalog().get(), Optional.empty()))
+                        .append(".");
+            }
+            builder.append(formatExpression(node.getSchema(), Optional.empty()));
+
+            return null;
+        }
+
         @Override
         protected Void visitShowSession(ShowSession node, Integer context)
         {

diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/tree/Use.java b/presto-parser/src/main/java/com/facebook/presto/sql/tree/Use.java
@@ -100,6 +100,9 @@ public boolean equals(Object o)
     @Override
     public String toString()
     {
-        return toStringHelper(this).toString();
+        return toStringHelper(this)
+                .add("catalog", catalog)
+                .add("schema", schema)
+                .toString();
     }
 }
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
@@ -140,6 +140,7 @@
 import com.facebook.presto.sql.tree.TransactionAccessMode;
 import com.facebook.presto.sql.tree.Union;
 import com.facebook.presto.sql.tree.Unnest;
+import com.facebook.presto.sql.tree.Use;
 import com.facebook.presto.sql.tree.Values;
 import com.facebook.presto.sql.tree.Window;
 import com.facebook.presto.sql.tree.With;
@@ -2441,6 +2442,16 @@ public void testNullTreatment()
                         ImmutableList.of(new Identifier("x"), new LongLiteral("1"))));
     }
 
+    @Test
+    public void testUse()
+    {
+        assertStatement("Use test_schema", new Use(Optional.empty(), identifier("test_schema")));
+        assertStatement("Use test_catalog.test_schema", new Use(Optional.of(identifier("test_catalog")), new Identifier("test_schema")));
+
+        assertStatement("Use \"test_schema\"", new Use(Optional.empty(), quotedIdentifier("test_schema")));
+        assertStatement("Use \"test_catalog\".\"test_schema\"", new Use(Optional.of(quotedIdentifier("test_catalog")), quotedIdentifier("test_schema")));
+    }
+
     private static void assertCast(String type)
     {
         assertCast(type, type);