Examples with TenantCallBack com.alibaba.druid.wall.WallConfig.TenantCallBack used on opensource projects

Example 1 with TenantCallBack

use of com.alibaba.druid.wall.WallConfig.TenantCallBack in project druid by alibaba.

the class WallVisitorUtils method checkInsertForMultiTenant.

private static void checkInsertForMultiTenant(WallVisitor visitor, SQLInsertInto x) {
    TenantCallBack tenantCallBack = visitor.getConfig().getTenantCallBack();
    String tenantTablePattern = visitor.getConfig().getTenantTablePattern();
    if (tenantCallBack == null && (tenantTablePattern == null || tenantTablePattern.length() == 0)) {
        return;
    }
    if (x == null) {
        throw new IllegalStateException("x is null");
    }
    SQLExprTableSource tableSource = x.getTableSource();
    String alias = null;
    String matchTableName = null;
    String tenantColumn = null;
    SQLExpr tableExpr = tableSource.getExpr();
    if (tableExpr instanceof SQLIdentifierExpr) {
        String tableName = ((SQLIdentifierExpr) tableExpr).getName();
        if (tenantCallBack != null) {
            tenantColumn = tenantCallBack.getTenantColumn(StatementType.INSERT, tableName);
        }
        if (StringUtils.isEmpty(tenantColumn) && ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName)) {
            tenantColumn = visitor.getConfig().getTenantColumn();
        }
        if (!StringUtils.isEmpty(tenantColumn)) {
            matchTableName = tableName;
            alias = tableSource.getAlias();
        }
    }
    if (matchTableName == null) {
        return;
    }
    SQLExpr item = null;
    if (alias != null) {
        item = new SQLPropertyExpr(new SQLIdentifierExpr(alias), tenantColumn);
    } else {
        item = new SQLIdentifierExpr(tenantColumn);
    }
    SQLExpr value = generateTenantValue(visitor, alias, StatementType.INSERT, matchTableName);
    // add insert item and value
    x.getColumns().add(item);
    List<ValuesClause> valuesClauses = null;
    ValuesClause valuesClause = null;
    if (x instanceof MySqlInsertStatement) {
        valuesClauses = ((MySqlInsertStatement) x).getValuesList();
    } else if (x instanceof SQLServerInsertStatement) {
        valuesClauses = ((MySqlInsertStatement) x).getValuesList();
    } else {
        valuesClause = x.getValues();
    }
    if (valuesClauses != null && valuesClauses.size() > 0) {
        for (ValuesClause clause : valuesClauses) {
            clause.addValue(value);
        }
    }
    if (valuesClause != null) {
        valuesClause.addValue(value);
    }
    // insert .. select
    SQLSelect select = x.getQuery();
    if (select != null) {
        List<SQLSelectQueryBlock> queryBlocks = splitSQLSelectQuery(select.getQuery());
        for (SQLSelectQueryBlock queryBlock : queryBlocks) {
            queryBlock.getSelectList().add(new SQLSelectItem(value));
        }
    }
    visitor.setSqlModified(true);
}

Example 2 with TenantCallBack

use of com.alibaba.druid.wall.WallConfig.TenantCallBack in project druid by alibaba.

the class WallVisitorUtils method checkUpdateForMultiTenant.

private static void checkUpdateForMultiTenant(WallVisitor visitor, SQLUpdateStatement x) {
    TenantCallBack tenantCallBack = visitor.getConfig().getTenantCallBack();
    String tenantTablePattern = visitor.getConfig().getTenantTablePattern();
    if (tenantCallBack == null && (tenantTablePattern == null || tenantTablePattern.length() == 0)) {
        return;
    }
    if (x == null) {
        throw new IllegalStateException("x is null");
    }
    SQLTableSource tableSource = x.getTableSource();
    String alias = null;
    String matchTableName = null;
    String tenantColumn = null;
    if (tableSource instanceof SQLExprTableSource) {
        SQLExpr tableExpr = ((SQLExprTableSource) tableSource).getExpr();
        if (tableExpr instanceof SQLIdentifierExpr) {
            String tableName = ((SQLIdentifierExpr) tableExpr).getName();
            if (tenantCallBack != null) {
                tenantColumn = tenantCallBack.getTenantColumn(StatementType.UPDATE, tableName);
            }
            if (StringUtils.isEmpty(tenantColumn) && ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName)) {
                tenantColumn = visitor.getConfig().getTenantColumn();
            }
            if (!StringUtils.isEmpty(tenantColumn)) {
                matchTableName = tableName;
                alias = tableSource.getAlias();
            }
        }
    }
    if (matchTableName == null) {
        return;
    }
    SQLExpr item = null;
    if (alias != null) {
        item = new SQLPropertyExpr(new SQLIdentifierExpr(alias), tenantColumn);
    } else {
        item = new SQLIdentifierExpr(tenantColumn);
    }
    SQLExpr value = generateTenantValue(visitor, alias, StatementType.UPDATE, matchTableName);
    SQLUpdateSetItem updateSetItem = new SQLUpdateSetItem();
    updateSetItem.setColumn(item);
    updateSetItem.setValue(value);
    x.addItem(updateSetItem);
    visitor.setSqlModified(true);
}

Example 3 with TenantCallBack

use of com.alibaba.druid.wall.WallConfig.TenantCallBack in project druid by alibaba.

the class WallVisitorUtils method checkSelectForMultiTenant.

private static void checkSelectForMultiTenant(WallVisitor visitor, SQLSelectQueryBlock x) {
    TenantCallBack tenantCallBack = visitor.getConfig().getTenantCallBack();
    String tenantTablePattern = visitor.getConfig().getTenantTablePattern();
    if (tenantCallBack == null && (tenantTablePattern == null || tenantTablePattern.length() == 0)) {
        return;
    }
    if (x == null) {
        throw new IllegalStateException("x is null");
    }
    if (!isSelectStatmentForMultiTenant(x)) {
        return;
    }
    SQLTableSource tableSource = x.getFrom();
    String alias = null;
    String matchTableName = null;
    String tenantColumn = null;
    if (tableSource instanceof SQLExprTableSource) {
        SQLExpr tableExpr = ((SQLExprTableSource) tableSource).getExpr();
        if (tableExpr instanceof SQLIdentifierExpr) {
            String tableName = ((SQLIdentifierExpr) tableExpr).getName();
            if (tenantCallBack != null) {
                tenantColumn = tenantCallBack.getTenantColumn(StatementType.SELECT, tableName);
            }
            if (StringUtils.isEmpty(tenantColumn) && ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName)) {
                tenantColumn = visitor.getConfig().getTenantColumn();
            }
            if (!StringUtils.isEmpty(tenantColumn)) {
                matchTableName = tableName;
                alias = tableSource.getAlias();
            }
        }
    } else if (tableSource instanceof SQLJoinTableSource) {
        SQLJoinTableSource join = (SQLJoinTableSource) tableSource;
        if (join.getLeft() instanceof SQLExprTableSource) {
            SQLExpr tableExpr = ((SQLExprTableSource) join.getLeft()).getExpr();
            if (tableExpr instanceof SQLIdentifierExpr) {
                String tableName = ((SQLIdentifierExpr) tableExpr).getName();
                if (tenantCallBack != null) {
                    tenantColumn = tenantCallBack.getTenantColumn(StatementType.SELECT, tableName);
                }
                if (StringUtils.isEmpty(tenantColumn) && ServletPathMatcher.getInstance().matches(tenantTablePattern, tableName)) {
                    tenantColumn = visitor.getConfig().getTenantColumn();
                }
                if (!StringUtils.isEmpty(tenantColumn)) {
                    matchTableName = tableName;
                    alias = join.getLeft().getAlias();
                    if (alias == null) {
                        alias = tableName;
                    }
                }
            }
            checkJoinSelectForMultiTenant(visitor, join, x);
        } else {
            checkJoinSelectForMultiTenant(visitor, join, x);
        }
    }
    if (matchTableName == null) {
        return;
    }
    SQLExpr item = null;
    if (alias != null) {
        item = new SQLPropertyExpr(new SQLIdentifierExpr(alias), tenantColumn);
    } else {
        item = new SQLIdentifierExpr(tenantColumn);
    }
    SQLSelectItem selectItem = new SQLSelectItem(item);
    x.getSelectList().add(selectItem);
    visitor.setSqlModified(true);
}

Example 4 with TenantCallBack

use of com.alibaba.druid.wall.WallConfig.TenantCallBack in project druid by alibaba.

the class WallFilter method preprocessResultSet.

private void preprocessResultSet(ResultSetProxy resultSet) throws SQLException {
    if (resultSet == null) {
        return;
    }
    ResultSetMetaData metaData = resultSet.getResultSetRaw().getMetaData();
    if (metaData == null) {
        return;
    }
    TenantCallBack tenantCallBack = provider.getConfig().getTenantCallBack();
    String tenantTablePattern = provider.getConfig().getTenantTablePattern();
    if (tenantCallBack == null && (tenantTablePattern == null || tenantTablePattern.length() == 0)) {
        return;
    }
    Map<Integer, Integer> logicColumnMap = new HashMap<Integer, Integer>();
    Map<Integer, Integer> physicalColumnMap = new HashMap<Integer, Integer>();
    List<Integer> hiddenColumns = new ArrayList<Integer>();
    List<Integer> tenantColumns = new ArrayList<Integer>();
    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) || 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)) {
            String columnName = metaData.getColumnName(physicalColumn);
            if (null != hiddenColumn && hiddenColumn.equalsIgnoreCase(columnName)) {
                hiddenColumns.add(physicalColumn);
                isHidden = true;
            }
        }
        if (!isHidden) {
            logicColumnMap.put(logicColumn, physicalColumn);
            physicalColumnMap.put(physicalColumn, logicColumn);
            logicColumn++;
        }
        if (!StringUtils.isEmpty(tenantColumn) && null != tenantColumn && tenantColumn.equalsIgnoreCase(metaData.getColumnName(physicalColumn))) {
            tenantColumns.add(physicalColumn);
        }
    }
    if (hiddenColumns.size() > 0) {
        resultSet.setLogicColumnMap(logicColumnMap);
        resultSet.setPhysicalColumnMap(physicalColumnMap);
        resultSet.setHiddenColumns(hiddenColumns);
    }
    tenantColumnsLocal.set(tenantColumns);
}

Example 5 with TenantCallBack

use of com.alibaba.druid.wall.WallConfig.TenantCallBack in project druid by alibaba.

the class WallVisitorUtils method generateTenantValue.

private static SQLExpr generateTenantValue(WallVisitor visitor, String alias, StatementType statementType, String tableName) {
    SQLExpr value;
    TenantCallBack callBack = visitor.getConfig().getTenantCallBack();
    if (callBack != null) {
        WallProvider.setTenantValue(callBack.getTenantValue(statementType, tableName));
    }
    Object tenantValue = WallProvider.getTenantValue();
    if (tenantValue instanceof Number) {
        value = new SQLNumberExpr((Number) tenantValue);
    } else if (tenantValue instanceof String) {
        value = new SQLCharExpr((String) tenantValue);
    } else {
        throw new IllegalStateException("tenant value not support type " + tenantValue);
    }
    return value;
}