diff --git a/src/main/java/com/alibaba/druid/sql/ast/expr/SQLInSubQueryExpr.java b/src/main/java/com/alibaba/druid/sql/ast/expr/SQLInSubQueryExpr.java index 175e4ff594..88ad058ccc 100644 --- a/src/main/java/com/alibaba/druid/sql/ast/expr/SQLInSubQueryExpr.java +++ b/src/main/java/com/alibaba/druid/sql/ast/expr/SQLInSubQueryExpr.java @@ -70,6 +70,7 @@ public class SQLInSubQueryExpr extends SQLExprImpl implements Serializable { @Override protected void accept0(SQLASTVisitor visitor) { if (visitor.visit(this)) { + acceptChild(visitor,this.expr); acceptChild(visitor, this.subQuery); } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGCurrentOfExpr.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGCurrentOfExpr.java index b8a2dd16ac..0d1d484232 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGCurrentOfExpr.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGCurrentOfExpr.java @@ -7,6 +7,14 @@ public class PGCurrentOfExpr extends PGSQLObjectImpl implements SQLExpr { private static final long serialVersionUID = 1L; private SQLExpr cursor; + + public PGCurrentOfExpr() { + + } + + public PGCurrentOfExpr(SQLExpr cursor) { + this.cursor = cursor; + } public SQLExpr getCursor() { return cursor; diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGDeleteStatement.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGDeleteStatement.java index 451d7512f3..ba9c98a610 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGDeleteStatement.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/ast/PGDeleteStatement.java @@ -14,8 +14,17 @@ public class PGDeleteStatement extends SQLDeleteStatement implements PGSQLStatem private PGWithClause with; private boolean only = false; private List using = new ArrayList(2); + private boolean returning; private String alias; + public boolean isReturning() { + return returning; + } + + public void setReturning(boolean returning) { + this.returning = returning; + } + public String getAlias() { return alias; } @@ -47,7 +56,7 @@ public class PGDeleteStatement extends SQLDeleteStatement implements PGSQLStatem public void setWith(PGWithClause with) { this.with = with; } - + protected void accept0(SQLASTVisitor visitor) { accept0((PGASTVisitor) visitor); } diff --git a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java index 2860728409..e09362fb13 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/postgresql/parser/PGLexer.java @@ -105,6 +105,8 @@ public class PGLexer extends Lexer { map.put("CASCADE", Token.CASCADE); map.put("RESTRICT", Token.RESTRICT); map.put("USING", Token.USING); + map.put("CURRENT", Token.CURRENT); + map.put("RETURNING", Token.RETURNING); DEFAULT_MYSQL_KEYWORDS = new Keywords(map); } 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 5793767f52..b68065d45a 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 @@ -6,6 +6,7 @@ import com.alibaba.druid.sql.ast.SQLExpr; import com.alibaba.druid.sql.ast.SQLName; import com.alibaba.druid.sql.ast.SQLStatement; import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; +import com.alibaba.druid.sql.dialect.postgresql.ast.PGCurrentOfExpr; import com.alibaba.druid.sql.dialect.postgresql.ast.PGDeleteStatement; import com.alibaba.druid.sql.dialect.postgresql.ast.PGTruncateStatement; import com.alibaba.druid.sql.parser.Lexer; @@ -59,8 +60,23 @@ public class PGSQLStatementParser extends SQLStatementParser { if (lexer.token() == (Token.WHERE)) { lexer.nextToken(); - SQLExpr where = this.exprParser.expr(); - deleteStatement.setWhere(where); + + if (lexer.token() == Token.CURRENT) { + lexer.nextToken(); + accept(Token.OF); + SQLName cursorName = this.exprParser.name(); + SQLExpr where = new PGCurrentOfExpr(cursorName); + deleteStatement.setWhere(where); + } else { + SQLExpr where = this.exprParser.expr(); + deleteStatement.setWhere(where); + } + } + + if (lexer.token() == Token.RETURNING) { + lexer.nextToken(); + accept(Token.STAR); + deleteStatement.setReturning(true); } return deleteStatement; 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 378235a993..a64849c3ea 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 @@ -303,6 +303,11 @@ public class PGOutputVisitor extends SQLASTOutputVisitor implements PGASTVisitor print("WHERE "); x.getWhere().accept(this); } + + if (x.isReturning()) { + println(); + print("RETURNING *"); + } return false; } diff --git a/src/main/java/com/alibaba/druid/sql/parser/Token.java b/src/main/java/com/alibaba/druid/sql/parser/Token.java index 04d4f20a85..01551c75c0 100644 --- a/src/main/java/com/alibaba/druid/sql/parser/Token.java +++ b/src/main/java/com/alibaba/druid/sql/parser/Token.java @@ -125,6 +125,8 @@ public enum Token { CASCADE("CASCADE"), RESTRICT("RESTRICT"), USING("USING"), + CURRENT("CURRENT"), + RETURNING("RETURNING"), // oracle START("START"), diff --git a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java index 6f0e3f9617..36a725e446 100644 --- a/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java +++ b/src/main/java/com/alibaba/druid/sql/visitor/SchemaStatVisitor.java @@ -11,6 +11,7 @@ import com.alibaba.druid.sql.ast.SQLObject; import com.alibaba.druid.sql.ast.expr.SQLAggregateExpr; import com.alibaba.druid.sql.ast.expr.SQLAllColumnExpr; import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr; +import com.alibaba.druid.sql.ast.expr.SQLInSubQueryExpr; import com.alibaba.druid.sql.ast.expr.SQLMethodInvokeExpr; import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr; import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement; @@ -364,6 +365,11 @@ public class SchemaStatVisitor extends SQLASTVisitorAdapter { return false; } + + @Override + public boolean visit(SQLInSubQueryExpr x) { + return true; + } public void endVisit(SQLDeleteStatement x) { aliasLocal.set(null); diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java index c176044aaa..a5ace197be 100644 --- a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest2.java @@ -30,7 +30,7 @@ public class PGDeleteTest2 extends PGTest { Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("tasks"))); - //Assert.assertTrue(visitor.getFields().size() == 0); + Assert.assertTrue(visitor.getFields().size() == 0); } diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java new file mode 100644 index 0000000000..d49397e8f0 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest3.java @@ -0,0 +1,37 @@ +package com.alibaba.druid.bvt.sql.postgresql; + +import java.util.List; + +import junit.framework.Assert; + +import com.alibaba.druid.sql.PGTest; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.dialect.postgresql.parser.PGSQLStatementParser; +import com.alibaba.druid.sql.dialect.postgresql.visitor.PGSchemaStatVisitor; +import com.alibaba.druid.stat.TableStat; + +public class PGDeleteTest3 extends PGTest { + + public void test_0() throws Exception { + String sql = "DELETE FROM tasks WHERE status = 'DONE' RETURNING *;"; + + PGSQLStatementParser parser = new PGSQLStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement statemen = statementList.get(0); + System.out.println(output(statementList)); + + Assert.assertEquals(1, statementList.size()); + + PGSchemaStatVisitor visitor = new PGSchemaStatVisitor(); + statemen.accept(visitor); + + System.out.println("Tables : " + visitor.getTables()); + System.out.println("fields : " + visitor.getFields()); + + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("tasks"))); + + Assert.assertTrue(visitor.getFields().size() == 1); + } + + +} diff --git a/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java new file mode 100644 index 0000000000..efa14b7516 --- /dev/null +++ b/src/test/java/com/alibaba/druid/bvt/sql/postgresql/PGDeleteTest4.java @@ -0,0 +1,42 @@ +package com.alibaba.druid.bvt.sql.postgresql; + +import java.util.List; + +import junit.framework.Assert; + +import com.alibaba.druid.sql.PGTest; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.dialect.postgresql.parser.PGSQLStatementParser; +import com.alibaba.druid.sql.dialect.postgresql.visitor.PGSchemaStatVisitor; +import com.alibaba.druid.stat.TableStat; + +public class PGDeleteTest4 extends PGTest { + + public void test_0() throws Exception { + String sql = "DELETE FROM films WHERE producer_id IN (SELECT id FROM producers WHERE name = 'foo');"; + + PGSQLStatementParser parser = new PGSQLStatementParser(sql); + List statementList = parser.parseStatementList(); + SQLStatement statemen = statementList.get(0); + System.out.println(output(statementList)); + + Assert.assertEquals(1, statementList.size()); + + PGSchemaStatVisitor visitor = new PGSchemaStatVisitor(); + statemen.accept(visitor); + + System.out.println("Tables : " + visitor.getTables()); + System.out.println("fields : " + visitor.getFields()); + + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("films"))); + Assert.assertTrue(visitor.getTables().containsKey(new TableStat.Name("producers"))); + + Assert.assertTrue(visitor.getFields().size() == 3); + + Assert.assertTrue(visitor.getFields().contains(new TableStat.Column("films", "producer_id"))); + Assert.assertTrue(visitor.getFields().contains(new TableStat.Column("producers", "id"))); + Assert.assertTrue(visitor.getFields().contains(new TableStat.Column("producers", "name"))); + } + + +}