From 91fc66855349bbfc0103749511fb81e2772f0deb Mon Sep 17 00:00:00 2001
From: wenshao <szujobs@hotmail.com>
Date: Sat, 3 Aug 2013 09:27:52 +0800
Subject: [PATCH] improve sql-wall

---
 .../druid/wall/spi/DB2WallVisitor.java        | 19 +++++++++++
 .../druid/wall/spi/OracleWallVisitor.java     | 32 ++++++++++++++++++-
 .../alibaba/druid/wall/spi/PGWallVisitor.java | 20 ++++++++++++
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java
index 6d44b49950..6a051ca40e 100644
--- a/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java
+++ b/src/main/java/com/alibaba/druid/wall/spi/DB2WallVisitor.java
@@ -187,16 +187,29 @@ public class DB2WallVisitor extends DB2ASTVisitorAdapter implements WallVisitor
                                                                    this.toSQL(x)));
             return false;
         }
+        
+        WallVisitorUtils.initWallTopStatementContext();
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLSelectStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLInsertStatement x) {
+        WallVisitorUtils.initWallTopStatementContext();
         WallVisitorUtils.checkInsert(this, x);
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLInsertStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLDeleteStatement x) {
@@ -206,10 +219,16 @@ public class DB2WallVisitor extends DB2ASTVisitorAdapter implements WallVisitor
 
     @Override
     public boolean visit(SQLUpdateStatement x) {
+        WallVisitorUtils.initWallTopStatementContext();
         WallVisitorUtils.checkUpdate(this, x);
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLUpdateStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLSelectItem x) {
diff --git a/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java
index 6897b64b60..6703a6b91c 100644
--- a/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java
+++ b/src/main/java/com/alibaba/druid/wall/spi/OracleWallVisitor.java
@@ -68,7 +68,7 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
         this.config = provider.getConfig();
         this.provider = provider;
     }
-    
+
     @Override
     public boolean isSqlModified() {
         return sqlModified;
@@ -204,6 +204,11 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLSelectStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(OracleInsertStatement x) {
@@ -218,6 +223,16 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
         return true;
     }
 
+    @Override
+    public void endVisit(OracleInsertStatement x) {
+        endVisit((SQLInsertStatement) x);
+    }
+
+    @Override
+    public void endVisit(SQLInsertStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
+
     @Override
     public boolean visit(InsertIntoClause x) {
         WallVisitorUtils.checkInsert(this, x);
@@ -237,6 +252,11 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
         return true;
     }
 
+    @Override
+    public void endVisit(OracleMultiInsertStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
+
     @Override
     public boolean visit(OracleDeleteStatement x) {
         return visit((SQLDeleteStatement) x);
@@ -247,6 +267,16 @@ public class OracleWallVisitor extends OracleASTVisitorAdapter implements WallVi
         WallVisitorUtils.checkDelete(this, x);
         return true;
     }
+    
+    @Override
+    public void endVisit(OracleDeleteStatement x) {
+        endVisit((SQLDeleteStatement) x);
+    }
+
+    @Override
+    public void endVisit(SQLDeleteStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(OracleUpdateStatement x) {
diff --git a/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java b/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java
index 478b87a51f..2823cf0391 100644
--- a/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java
+++ b/src/main/java/com/alibaba/druid/wall/spi/PGWallVisitor.java
@@ -193,6 +193,11 @@ public class PGWallVisitor extends PGASTVisitorAdapter implements WallVisitor {
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLSelectStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLInsertStatement x) {
@@ -201,12 +206,22 @@ public class PGWallVisitor extends PGASTVisitorAdapter implements WallVisitor {
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLInsertStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLDeleteStatement x) {
         WallVisitorUtils.checkDelete(this, x);
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLDeleteStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLUpdateStatement x) {
@@ -215,6 +230,11 @@ public class PGWallVisitor extends PGASTVisitorAdapter implements WallVisitor {
 
         return true;
     }
+    
+    @Override
+    public void endVisit(SQLUpdateStatement x) {
+        WallVisitorUtils.clearWallTopStatementContext();
+    }
 
     @Override
     public boolean visit(SQLSelectItem x) {