优化postgresql解析set schema和set names的逻辑 alibaba#5870

lizongbo · May 4, 2024 · 2496258 · 2496258
1 parent 0c432cb
commit 2496258
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 2 deletions.
diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java b/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java
@@ -23,6 +23,7 @@
 import com.alibaba.druid.sql.dialect.postgresql.ast.stmt.*;
 import com.alibaba.druid.sql.parser.*;
 import com.alibaba.druid.util.FnvHash;
+import com.alibaba.druid.util.JdbcUtils;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -667,6 +668,12 @@ public SQLStatement parseSet() {
             lexer.nextToken();
             values.add(this.exprParser.primary());
             lexer.nextToken();
+        } else if (JdbcUtils.isPgsqlDbType(dbType) && ("schema".equalsIgnoreCase(parameter) || "names".equalsIgnoreCase(parameter))) {
+            paramExpr = new SQLIdentifierExpr(parameter);
+            lexer.nextToken();
+            String value = lexer.stringVal();
+            values.add(new SQLCharExpr(value));
+            lexer.nextToken();
         } else {
             paramExpr = new SQLIdentifierExpr(parameter);
             lexer.nextToken();

diff --git a/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java b/core/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
@@ -653,8 +653,12 @@ public boolean visit(SQLAssignItem x) {
 
         x.getTarget().accept(this);
         SQLExpr value = x.getValue();
-        if (x.getTarget() instanceof SQLIdentifierExpr
-            && ((SQLIdentifierExpr) x.getTarget()).getName().equalsIgnoreCase("TIME ZONE")) {
+        boolean needSpace = false;
+        if (x.getTarget() instanceof SQLIdentifierExpr) {
+            String name = ((SQLIdentifierExpr) x.getTarget()).getName();
+            needSpace = "TIME ZONE".equalsIgnoreCase(name) || "schema".equalsIgnoreCase(name) || "names".equalsIgnoreCase(name);
+        }
+        if (needSpace) {
             print(' ');
         } else {
             if (!((SQLSetStatement) x.getParent()).isUseSet()) {

diff --git a/core/src/test/java/com/alibaba/druid/bvt/sql/postgresql/issues/Issue5870.java b/core/src/test/java/com/alibaba/druid/bvt/sql/postgresql/issues/Issue5870.java
@@ -0,0 +1,48 @@
+package com.alibaba.druid.bvt.sql.postgresql.issues;
+
+import java.util.List;
+
+import com.alibaba.druid.DbType;
+import com.alibaba.druid.sql.ast.SQLStatement;
+import com.alibaba.druid.sql.parser.SQLParserUtils;
+import com.alibaba.druid.sql.parser.SQLStatementParser;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * PostgreSQL解析SET SCHEMA的问题
+ *
+ * @author lizongbo
+ * @see <a href="https://github.com/alibaba/druid/issues/5870">Issue来源</a>
+ * @see <a href="https://www.postgresql.org/docs/current/sql-set.html">SET SCHEMA</a>
+ * @see <a href="https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-SETTING">SET NAMES 'value';</a>
+ */
+public class Issue5870 {
+
+    @Test
+    public void test_parse_set_schema() {
+        for (String sql : new String[]{
+            "SET SCHEMA 'platform_base';",
+            "SET search_path TO my_schema, public;",
+            "show client_encoding;",
+            "set client_encoding = GBK;",
+            "set client_encoding to 'UTF-8'",
+            "SET CLIENT_ENCODING TO 'value';",
+            "SET NAMES 'UTF-8'",
+        }) {
+            SQLStatementParser parser1 = SQLParserUtils.createSQLStatementParser(sql, DbType.postgresql);
+            List<SQLStatement> statementList1 = parser1.parseStatementList();
+            System.out.println("原始的sql===" + sql);
+            String sqleNew = statementList1.get(0).toString();
+            System.out.println("生成的sql===" + sqleNew);
+            SQLStatementParser parser2 = SQLParserUtils.createSQLStatementParser(sqleNew, DbType.postgresql);
+            List<SQLStatement> statementList2 = parser2.parseStatementList();
+            String sqleNew2 = statementList2.get(0).toString();
+            System.out.println("再次解析生成的sql===" + sqleNew);
+            assertEquals(sqleNew, sqleNew2);
+        }
+
+    }
+}