Skip to content

Commit 14c86ac

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

2 files changed

Lines changed: 187 additions & 0 deletions

File tree

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

Lines changed: 110 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,85 @@ 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+
// An analyzed recursive member consumes the current CTE through a non-anchor CTERelation.
840+
// Do not expand its definition again, or the collector will revisit the same recursive member.
841+
return null;
842+
}
843+
if (cteNameScopes.isEmpty()) {
844+
return super.visitCTE(node, context);
845+
}
846+
if (isInRecursiveCteScope() && node.getCteQueryStatement().getQueryRelation()
847+
instanceof SetOperationRelation setOperationRelation) {
848+
visitRecursiveCteDefinition(node, setOperationRelation, context);
849+
} else {
850+
visit(node.getCteQueryStatement(), context);
851+
}
852+
addCteName(cteNameScopes.peek(), node);
853+
return null;
854+
}
855+
856+
private void visitRecursiveCteDefinition(CTERelation cteRelation, SetOperationRelation setOperationRelation,
857+
Void context) {
858+
List<QueryRelation> relations = setOperationRelation.getRelations();
859+
if (relations.isEmpty()) {
860+
return;
861+
}
862+
// Mirror QueryAnalyzer.tryProcessRecursiveCte: the first set-op child is the anchor.
863+
// The current CTE name is not visible there, so a same-named table must still be collected.
864+
visit(relations.get(0), context);
865+
866+
// Only recursive members can see the current CTE name. Keep that visibility in a temporary
867+
// scope so unresolved self-references are skipped without hiding anchor base tables.
868+
Set<String> recursiveNames = new HashSet<>();
869+
addCteName(recursiveNames, cteRelation);
870+
cteNameScopes.push(recursiveNames);
871+
try {
872+
for (int i = 1; i < relations.size(); i++) {
873+
visit(relations.get(i), context);
874+
}
875+
} finally {
876+
cteNameScopes.pop();
877+
}
878+
}
879+
880+
private boolean isInRecursiveCteScope() {
881+
return !recursiveCteScopes.isEmpty() && recursiveCteScopes.peek();
882+
}
883+
884+
private void addCteName(Set<String> names, CTERelation cteRelation) {
885+
if (!Strings.isNullOrEmpty(cteRelation.getName())) {
886+
names.add(cteRelation.getName());
887+
}
888+
}
889+
890+
private boolean isUnresolvedCteReference(TableRelation node) {
891+
if (node.getTable() != null || cteNameScopes.isEmpty()) {
892+
return false;
893+
}
894+
TableName tableName = node.getName();
895+
if (tableName == null || !Strings.isNullOrEmpty(tableName.getCatalog()) ||
896+
!Strings.isNullOrEmpty(tableName.getDb()) || Strings.isNullOrEmpty(tableName.getTbl())) {
897+
return false;
898+
}
899+
String tableNameWithoutDb = tableName.getTbl();
900+
for (Set<String> cteNames : cteNameScopes) {
901+
if (cteNames.contains(tableNameWithoutDb)) {
902+
return true;
903+
}
904+
}
905+
return false;
906+
}
797907
}
798908

799909
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)