diff --git a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java index 22a6555fd9..2359790499 100644 --- a/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java +++ b/src/main/java/com/alibaba/druid/sql/dialect/mysql/parser/MySqlExprParser.java @@ -728,12 +728,6 @@ public class MySqlExprParser extends SQLExprParser { MySqlUnique unique = new MySqlUnique(); - if (identifierEquals("USING")) { - lexer.nextToken(); - unique.setIndexType(lexer.stringVal()); - lexer.nextToken(); - } - if (lexer.token() != Token.LPAREN) { SQLName name = name(); unique.setName(name); @@ -750,6 +744,12 @@ public class MySqlExprParser extends SQLExprParser { } accept(Token.RPAREN); + if (identifierEquals("USING")) { + lexer.nextToken(); + unique.setIndexType(lexer.stringVal()); + lexer.nextToken(); + } + return unique; } diff --git a/src/main/java/com/alibaba/druid/wall/WallConfig.java b/src/main/java/com/alibaba/druid/wall/WallConfig.java index 942e444367..07438d29be 100644 --- a/src/main/java/com/alibaba/druid/wall/WallConfig.java +++ b/src/main/java/com/alibaba/druid/wall/WallConfig.java @@ -723,11 +723,11 @@ public class WallConfig implements WallConfigMBean { String getHiddenColumn(String tableName); /** - * resultset返回值中如果包含hiddenColumn的回调函数 + * resultset返回值中如果包含tenantColumn的回调函数 * - * @param value hiddenColumn对应的值 + * @param value tenantColumn对应的值 */ - void resultset_hiddenColumn(Object value); + void filterResultsetTenantColumn(Object value); } } diff --git a/src/main/java/com/alibaba/druid/wall/WallFilter.java b/src/main/java/com/alibaba/druid/wall/WallFilter.java index 19230807fd..7e65bcf936 100644 --- a/src/main/java/com/alibaba/druid/wall/WallFilter.java +++ b/src/main/java/com/alibaba/druid/wall/WallFilter.java @@ -55,6 +55,7 @@ import com.alibaba.druid.util.JdbcUtils; import com.alibaba.druid.util.ServletPathMatcher; import com.alibaba.druid.util.StringUtils; import com.alibaba.druid.wall.WallConfig.TenantCallBack; +import com.alibaba.druid.wall.WallConfig.TenantCallBack.StatementType; import com.alibaba.druid.wall.spi.DB2WallProvider; import com.alibaba.druid.wall.spi.MySqlWallProvider; import com.alibaba.druid.wall.spi.OracleWallProvider; @@ -1214,11 +1215,11 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { boolean hasNext = chain.resultSet_next(resultSet); TenantCallBack callback = provider.getConfig().getTenantCallBack(); if (callback != null && hasNext) { - List hiddenColumns = resultSet.getHiddenColumns(); - if (hiddenColumns != null && hiddenColumns.size() > 0) { - for (Integer columnIndex : hiddenColumns) { + List tenantColumns = tenantColumnsLocal.get(); + if (tenantColumns != null && tenantColumns.size() > 0) { + for (Integer columnIndex : tenantColumns) { Object value = resultSet.getResultSetRaw().getObject(columnIndex); - callback.resultset_hiddenColumn(value); + callback.filterResultsetTenantColumn(value); } } } @@ -1371,6 +1372,8 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { return provider.checkValid(sql); } + private static final ThreadLocal> tenantColumnsLocal = new ThreadLocal>(); + private void preprocessResultSet(ResultSetProxy resultSet) throws SQLException { if (resultSet == null) { return; @@ -1390,17 +1393,28 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { Map logicColumnMap = new HashMap(); Map physicalColumnMap = new HashMap(); List hiddenColumns = new ArrayList(); + List tenantColumns = new ArrayList(); for (int physicalColumn = 1, logicColumn = 1; physicalColumn <= metaData.getColumnCount(); physicalColumn++) { boolean isHidden = false; String tableName = metaData.getTableName(physicalColumn); String hiddenColumn = null; + String tenantColumn = null; if (tenantCallBack != null) { + tenantColumn = tenantCallBack.getTenantColumn(StatementType.SELECT, tableName); hiddenColumn = tenantCallBack.getHiddenColumn(tableName); } - if (StringUtils.isEmpty(hiddenColumn) - && (tableName == null || ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName))) { - hiddenColumn = provider.getConfig().getTenantColumn(); + + if (StringUtils.isEmpty(hiddenColumn) || StringUtils.isEmpty(tenantColumn)) { + if (tableName == null || ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName)) { + if (StringUtils.isEmpty(hiddenColumn)) { + hiddenColumn = provider.getConfig().getTenantColumn(); + } + + if (StringUtils.isEmpty(tenantColumn)) { + tenantColumn = provider.getConfig().getTenantColumn(); + } + } } if (!StringUtils.isEmpty(hiddenColumn)) { @@ -1415,6 +1429,11 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { physicalColumnMap.put(physicalColumn, logicColumn); logicColumn++; } + + if (!StringUtils.isEmpty(tenantColumn) + && tenantColumn.equalsIgnoreCase(metaData.getColumnName(physicalColumn))) { + tenantColumns.add(physicalColumn); + } } if (hiddenColumns.size() > 0) { @@ -1422,5 +1441,6 @@ public class WallFilter extends FilterAdapter implements WallFilterMBean { resultSet.setPhysicalColumnMap(physicalColumnMap); resultSet.setHiddenColumns(hiddenColumns); } + tenantColumnsLocal.set(tenantColumns); } } diff --git a/src/test/java/com/alibaba/druid/bvt/filter/wall/TenantTestCallBack.java b/src/test/java/com/alibaba/druid/bvt/filter/wall/TenantTestCallBack.java index 7debb18557..6ddcc5750c 100644 --- a/src/test/java/com/alibaba/druid/bvt/filter/wall/TenantTestCallBack.java +++ b/src/test/java/com/alibaba/druid/bvt/filter/wall/TenantTestCallBack.java @@ -35,7 +35,7 @@ public class TenantTestCallBack implements TenantCallBack { } @Override - public void resultset_hiddenColumn(Object value) { + public void filterResultsetTenantColumn(Object value) { System.out.println(value); } }