From 97a1adcc46c1aa73d1b62eb2e72f05ac57e01af3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=AB=98=E9=93=81?= <shaojin.wensj@alibaba-inc.com>
Date: Tue, 3 Nov 2020 21:08:39 +0800
Subject: [PATCH] improved pg support, fix #4021

---
 .../statement/SQLCreateIndexStatement.java    | 10 ++++++++
 .../parser/PGSQLStatementParser.java          |  7 ++++++
 .../postgresql/visitor/PGOutputVisitor.java   |  4 +++
 .../sql/postgresql/PGCreateIndexTest_0.java   | 25 +++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateIndexStatement.java b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateIndexStatement.java
index 91dc93f1bc..cc2557fb61 100644
--- a/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateIndexStatement.java
+++ b/src/main/java/com/alibaba/druid/sql/ast/statement/SQLCreateIndexStatement.java
@@ -37,6 +37,7 @@ public class SQLCreateIndexStatement extends SQLStatementImpl implements SQLCrea
     protected List<SQLAssignItem> properties = new ArrayList<SQLAssignItem>();
     protected List<SQLAssignItem> tableProperties = new ArrayList<SQLAssignItem>();
     protected boolean storing;
+    protected boolean ifNotExists;
 
     public SQLCreateIndexStatement() {
         indexDefinition.setParent(this);
@@ -162,6 +163,7 @@ public class SQLCreateIndexStatement extends SQLStatementImpl implements SQLCrea
     public SQLCreateIndexStatement clone() {
         SQLCreateIndexStatement x = new SQLCreateIndexStatement();
         indexDefinition.cloneTo(x.indexDefinition);
+        x.setIfNotExists(ifNotExists);
         return x;
     }
 
@@ -318,4 +320,12 @@ public class SQLCreateIndexStatement extends SQLStatementImpl implements SQLCrea
     public List<SQLSelectOrderByItem> getColumns() {
         return indexDefinition.getColumns();
     }
+
+    public boolean isIfNotExists() {
+        return ifNotExists;
+    }
+
+    public void setIfNotExists(boolean ifNotExists) {
+        this.ifNotExists = ifNotExists;
+    }
 }
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java
index 81b3c3371c..c1892839c1 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGSQLStatementParser.java
@@ -608,6 +608,13 @@ public class PGSQLStatementParser extends SQLStatementParser {
 
         accept(Token.INDEX);
 
+        if (lexer.token() == Token.IF) {
+            lexer.nextToken();
+            accept(Token.NOT);
+            accept(Token.EXISTS);
+            stmt.setIfNotExists(true);
+        }
+
         if (lexer.identifierEquals(FnvHash.Constants.CONCURRENTLY)) {
             lexer.nextToken();
             stmt.setConcurrently(true);
diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
index 8d28b895ec..70d914cd62 100644
--- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
+++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/visitor/PGOutputVisitor.java
@@ -2401,6 +2401,10 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor
 
         print0(ucase ? "INDEX" : "index");
 
+        if (x.isIfNotExists()) {
+            print0(ucase ? " IF NOT EXISTS" : " if not exists");
+        }
+
         if (x.isConcurrently()) {
             print0(ucase ? " CONCURRENTLY" : " concurrently");
         }
diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGCreateIndexTest_0.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGCreateIndexTest_0.java
index c2207e09e6..d4d7e74fd3 100644
--- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGCreateIndexTest_0.java
+++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGCreateIndexTest_0.java
@@ -52,4 +52,29 @@ public class PGCreateIndexTest_0 extends PGTest {
         assertEquals(1, visitor.getColumns().size() );
     }
 
+    public void test_1() throws Exception {
+        String sql = "CREATE INDEX if not exists \"calculate_measure_dataset_id_idx\" ON \"calculate_measure\" USING btree (\"dataset_id\");";
+
+        PGSQLStatementParser parser = new PGSQLStatementParser(sql);
+        List<SQLStatement> statementList = parser.parseStatementList();
+        SQLStatement stmt = statementList.get(0);
+
+        assertEquals("CREATE INDEX IF NOT EXISTS \"calculate_measure_dataset_id_idx\" ON \"calculate_measure\" USING btree (\"dataset_id\");", SQLUtils.toPGString(stmt));
+
+        assertEquals("create index if not exists \"calculate_measure_dataset_id_idx\" on \"calculate_measure\" using btree (\"dataset_id\");", stmt.toLowerCaseString());
+
+        assertEquals(1, statementList.size());
+
+        PGSchemaStatVisitor visitor = new PGSchemaStatVisitor();
+        stmt.accept(visitor);
+
+//        System.out.println("Tables : " + visitor.getTables());
+//        System.out.println("fields : " + visitor.getColumns());
+
+//        assertTrue(visitor.getTables().containsKey(new TableStat.Name("public.city")));
+//
+//        assertTrue(visitor.getTables().get(new TableStat.Name("public.city")).getCreateIndexCount() == 1);
+//
+//        assertEquals(1, visitor.getColumns().size() );
+    }
 }