From 24962581e9ec7cc464f05816f1d4a3204ee41d98 Mon Sep 17 00:00:00 2001
From: LiZongbo <lizongbo@gmail.com>
Date: Sat, 4 May 2024 11:06:27 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96postgresql=E8=A7=A3=E6=9E=90s?=
 =?UTF-8?q?et=20schema=E5=92=8Cset=20names=E7=9A=84=E9=80=BB=E8=BE=91=20#5?=
 =?UTF-8?q?870?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

优化postgresql解析set schema和set names的逻辑 #5870
---
 .../parser/PGSQLStatementParser.java          |  7 +++
 .../postgresql/visitor/PGOutputVisitor.java   |  8 +++-
 .../bvt/sql/postgresql/issues/Issue5870.java  | 48 +++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 core/src/test/java/com/alibaba/druid/bvt/sql/postgresql/issues/Issue5870.java

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
index c0d945243f..c653df77b8 100644
--- 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
index 04be284848..98429c403a 100644
--- 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
new file mode 100644
index 0000000000..fe4959187b
--- /dev/null
+++ 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);
+        }
+
+    }
+}