Skip to content

Commit 74715a8

Browse files
HangyuanLiuHangyuanLiu
committed
Fix CTE table collection in security views
Signed-off-by: Harbor Liu <460660596@qq.com>
1 parent 6f25e14 commit 74715a8

2 files changed

Lines changed: 181 additions & 0 deletions

File tree

‎fe/fe-core/src/main/java/com/starrocks/sql/analyzer/AnalyzerUtils.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
import com.starrocks.sql.ast.PartitionKeyDesc;
8383
import com.starrocks.sql.ast.PartitionValue;
8484
import com.starrocks.sql.ast.QualifiedName;
85+
import com.starrocks.sql.ast.QueryRelation;
8586
import com.starrocks.sql.ast.QueryStatement;
8687
import com.starrocks.sql.ast.RangePartitionDesc;
8788
import com.starrocks.sql.ast.Relation;
@@ -140,9 +141,11 @@
140141

141142
import java.time.LocalDateTime;
142143
import java.time.format.DateTimeFormatter;
144+
import java.util.ArrayDeque;
143145
import java.util.ArrayList;
144146
import java.util.Arrays;
145147
import java.util.Collections;
148+
import java.util.Deque;
146149
import java.util.HashMap;
147150
import java.util.HashSet;
148151
import java.util.Iterator;
@@ -738,6 +741,8 @@ public static Map<TableName, Table> collectAllTable(StatementBase statementBase)
738741

739742
private static class TableCollector extends AstTraverser<Void, Void> {
740743
protected Map<TableName, Table> tables;
744+
private final Deque<Set<String>> cteNameScopes = new ArrayDeque<>();
745+
private final Deque<Boolean> recursiveCteScopes = new ArrayDeque<>();
741746

742747
public TableCollector() {
743748
this.tables = Maps.newHashMap();
@@ -754,6 +759,36 @@ public Void visitQueryStatement(QueryStatement statement, Void context) {
754759
return visit(statement.getQueryRelation());
755760
}
756761

762+
@Override
763+
public Void visitSelect(SelectRelation node, Void context) {
764+
if (!node.hasWithClause()) {
765+
return super.visitSelect(node, context);
766+
}
767+
cteNameScopes.push(new HashSet<>());
768+
recursiveCteScopes.push(node.isHasRecursiveCTE());
769+
try {
770+
return super.visitSelect(node, context);
771+
} finally {
772+
recursiveCteScopes.pop();
773+
cteNameScopes.pop();
774+
}
775+
}
776+
777+
@Override
778+
public Void visitSetOp(SetOperationRelation node, Void context) {
779+
if (!node.hasWithClause()) {
780+
return super.visitSetOp(node, context);
781+
}
782+
cteNameScopes.push(new HashSet<>());
783+
recursiveCteScopes.push(node.isHasRecursiveCTE());
784+
try {
785+
return super.visitSetOp(node, context);
786+
} finally {
787+
recursiveCteScopes.pop();
788+
cteNameScopes.pop();
789+
}
790+
}
791+
757792
// ------------------------------------------- DML Statement -------------------------------------------------------
758793

759794
@Override
@@ -790,10 +825,79 @@ public Void visitMergeIntoStatement(MergeIntoStmt node, Void context) {
790825

791826
@Override
792827
public Void visitTable(TableRelation node, Void context) {
828+
if (isUnresolvedCteReference(node)) {
829+
return null;
830+
}
793831
Table table = node.getTable();
794832
tables.put(node.getName(), table);
795833
return null;
796834
}
835+
836+
@Override
837+
public Void visitCTE(CTERelation node, Void context) {
838+
if (node.isRecursive() && !node.isAnchor()) {
839+
return null;
840+
}
841+
if (cteNameScopes.isEmpty()) {
842+
return super.visitCTE(node, context);
843+
}
844+
if (isInRecursiveCteScope() && node.getCteQueryStatement().getQueryRelation()
845+
instanceof SetOperationRelation setOperationRelation) {
846+
visitRecursiveCteDefinition(node, setOperationRelation, context);
847+
} else {
848+
visit(node.getCteQueryStatement(), context);
849+
}
850+
addCteName(cteNameScopes.peek(), node);
851+
return null;
852+
}
853+
854+
private void visitRecursiveCteDefinition(CTERelation cteRelation, SetOperationRelation setOperationRelation,
855+
Void context) {
856+
List<QueryRelation> relations = setOperationRelation.getRelations();
857+
if (relations.isEmpty()) {
858+
return;
859+
}
860+
visit(relations.get(0), context);
861+
862+
Set<String> recursiveNames = new HashSet<>();
863+
addCteName(recursiveNames, cteRelation);
864+
cteNameScopes.push(recursiveNames);
865+
try {
866+
for (int i = 1; i < relations.size(); i++) {
867+
visit(relations.get(i), context);
868+
}
869+
} finally {
870+
cteNameScopes.pop();
871+
}
872+
}
873+
874+
private boolean isInRecursiveCteScope() {
875+
return !recursiveCteScopes.isEmpty() && recursiveCteScopes.peek();
876+
}
877+
878+
private void addCteName(Set<String> names, CTERelation cteRelation) {
879+
if (!Strings.isNullOrEmpty(cteRelation.getName())) {
880+
names.add(cteRelation.getName());
881+
}
882+
}
883+
884+
private boolean isUnresolvedCteReference(TableRelation node) {
885+
if (node.getTable() != null || cteNameScopes.isEmpty()) {
886+
return false;
887+
}
888+
TableName tableName = node.getName();
889+
if (tableName == null || !Strings.isNullOrEmpty(tableName.getCatalog()) ||
890+
!Strings.isNullOrEmpty(tableName.getDb()) || Strings.isNullOrEmpty(tableName.getTbl())) {
891+
return false;
892+
}
893+
String tableNameWithoutDb = tableName.getTbl();
894+
for (Set<String> cteNames : cteNameScopes) {
895+
if (cteNames.contains(tableNameWithoutDb)) {
896+
return true;
897+
}
898+
}
899+
return false;
900+
}
797901
}
798902

799903
public static Map<TableName, Table> collectAllTableWithAlias(StatementBase statementBase) {

‎fe/fe-core/src/test/java/com/starrocks/sql/analyzer/AnalyzeUtilTest.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,14 @@
2626
import com.starrocks.qe.ConnectContext;
2727
import com.starrocks.server.GlobalStateMgr;
2828
import com.starrocks.sql.ast.AddPartitionClause;
29+
import com.starrocks.sql.ast.CTERelation;
2930
import com.starrocks.sql.ast.CreateViewStmt;
3031
import com.starrocks.sql.ast.ListPartitionDesc;
3132
import com.starrocks.sql.ast.MultiItemListPartitionDesc;
3233
import com.starrocks.sql.ast.PartitionDesc;
34+
import com.starrocks.sql.ast.QueryStatement;
35+
import com.starrocks.sql.ast.SelectRelation;
36+
import com.starrocks.sql.ast.SetOperationRelation;
3337
import com.starrocks.sql.ast.StatementBase;
3438
import com.starrocks.sql.parser.SqlParser;
3539
import com.starrocks.utframe.UtFrameUtils;
@@ -175,6 +179,79 @@ public void testCollectTable() {
175179
Assertions.assertEquals("[test.tprimary2, test.tprimary]", m.keySet().toString());
176180
}
177181

182+
@Test
183+
public void testCollectTableAndViewSkipsCteReferenceWithoutAnalyze() {
184+
String sql = "with x as (select v1 from db2.t0) " +
185+
"select db1.t0.v1 from db1.t0 join x on db1.t0.v1 = x.v1";
186+
StatementBase statementBase = SqlParser.parse(sql, 0L).get(0);
187+
Map<TableName, Table> m = AnalyzerUtils.collectAllTableAndView(statementBase);
188+
Set<String> tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
189+
Set<String> expectedTables = new HashSet<>(Lists.newArrayList("db1.t0", "db2.t0"));
190+
191+
Assertions.assertEquals(expectedTables, tableNames);
192+
}
193+
194+
@Test
195+
public void testCollectTableAndViewKeepsBaseTableInsideCteDefinitionWithoutAnalyze() {
196+
String sql = "with t0 as (select * from t0) select * from t0";
197+
StatementBase statementBase = SqlParser.parse(sql, 0L).get(0);
198+
Map<TableName, Table> m = AnalyzerUtils.collectAllTableAndView(statementBase);
199+
Set<String> tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
200+
Set<String> expectedTables = new HashSet<>(Lists.newArrayList("t0"));
201+
202+
Assertions.assertEquals(expectedTables, tableNames);
203+
204+
sql = "with x as (select * from y), y as (select * from t0) select * from y";
205+
statementBase = SqlParser.parse(sql, 0L).get(0);
206+
m = AnalyzerUtils.collectAllTableAndView(statementBase);
207+
tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
208+
expectedTables = new HashSet<>(Lists.newArrayList("y", "t0"));
209+
210+
Assertions.assertEquals(expectedTables, tableNames);
211+
212+
sql = "with X as (select * from t0) select * from x";
213+
statementBase = SqlParser.parse(sql, 0L).get(0);
214+
m = AnalyzerUtils.collectAllTableAndView(statementBase);
215+
tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
216+
expectedTables = new HashSet<>(Lists.newArrayList("t0", "x"));
217+
218+
Assertions.assertEquals(expectedTables, tableNames);
219+
220+
sql = "with recursive cte(n) as " +
221+
"(select v1 from t0 union all select n from cte where n < 10) select * from cte";
222+
statementBase = SqlParser.parse(sql, 0L).get(0);
223+
m = AnalyzerUtils.collectAllTableAndView(statementBase);
224+
tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
225+
expectedTables = new HashSet<>(Lists.newArrayList("t0"));
226+
227+
Assertions.assertEquals(expectedTables, tableNames);
228+
229+
sql = "with recursive cte(n) as " +
230+
"(select v1 from cte union all select n from cte where n < 10) select * from cte";
231+
statementBase = SqlParser.parse(sql, 0L).get(0);
232+
m = AnalyzerUtils.collectAllTableAndView(statementBase);
233+
tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
234+
expectedTables = new HashSet<>(Lists.newArrayList("cte"));
235+
236+
Assertions.assertEquals(expectedTables, tableNames);
237+
238+
sql = "with recursive cte as " +
239+
"(select v1 from t0 union all select v1 + 1 from cte where v1 < 10) select * from cte";
240+
statementBase = SqlParser.parse(sql, 0L).get(0);
241+
QueryStatement queryStatement = (QueryStatement) statementBase;
242+
CTERelation cteRelation = queryStatement.getQueryRelation().getCteRelations().get(0);
243+
SetOperationRelation setOperationRelation =
244+
(SetOperationRelation) cteRelation.getCteQueryStatement().getQueryRelation();
245+
SelectRelation recursiveMember = (SelectRelation) setOperationRelation.getRelations().get(1);
246+
recursiveMember.setRelation(new CTERelation(cteRelation.getCteMouldId(), cteRelation.getName(),
247+
cteRelation.getColumnOutputNames(), cteRelation.getCteQueryStatement(), true, false));
248+
m = AnalyzerUtils.collectAllTableAndView(statementBase);
249+
tableNames = m.keySet().stream().map(TableName::toString).collect(Collectors.toSet());
250+
expectedTables = new HashSet<>(Lists.newArrayList("t0"));
251+
252+
Assertions.assertEquals(expectedTables, tableNames);
253+
}
254+
178255
@Test
179256
public void testQueryStatementCollectColumns() {
180257
// base

0 commit comments

Comments
 (0)