Skip to content

Commit ae19df1

Browse files
committed
Fix check_access_control_on_utilized_columns_only for CTE
Previously we weren't associating the with query with the columns used from it in the main part of the query, so we wouldn't consider any columns as used. When the usage within the cte or outside of it was just an identifier, we would have the utilized column collected anyway because the column collected in the main query kept the source table information. However, if there was an expression in between, we would lose that information, and we wouldn't check access control for the column that was used in an expression in the cte. Example: For the following query ``` with cte as (select x as c1, z + 1 as c2 from t13) select c1, c2 from (select * from cte) ``` Previously we would only check access permissions on t13.x, but not t13.z. With this change we will check column access on both 13.x and t13.z.
1 parent 56bf5e6 commit ae19df1

File tree

3 files changed

+37
-0
lines changed

3 files changed

+37
-0
lines changed

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/RelationId.java

+4
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ public boolean isAnonymous()
5353
return sourceNode == null;
5454
}
5555

56+
public Node getSourceNode()
57+
{
58+
return sourceNode;
59+
}
5660
@Override
5761
public boolean equals(Object o)
5862
{

presto-main/src/main/java/com/facebook/presto/sql/analyzer/UtilizedColumnsAnalyzer.java

+27
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import com.facebook.presto.sql.tree.Lateral;
3838
import com.facebook.presto.sql.tree.Node;
3939
import com.facebook.presto.sql.tree.NodeRef;
40+
import com.facebook.presto.sql.tree.QualifiedName;
4041
import com.facebook.presto.sql.tree.Query;
4142
import com.facebook.presto.sql.tree.QuerySpecification;
4243
import com.facebook.presto.sql.tree.Relation;
@@ -48,6 +49,7 @@
4849
import com.facebook.presto.sql.tree.Union;
4950
import com.facebook.presto.sql.tree.Unnest;
5051
import com.facebook.presto.sql.tree.Values;
52+
import com.facebook.presto.sql.tree.WithQuery;
5153
import com.google.common.collect.HashMultimap;
5254
import com.google.common.collect.ImmutableSet;
5355

@@ -57,6 +59,8 @@
5759
import java.util.Map.Entry;
5860
import java.util.Set;
5961

62+
import static com.google.common.base.Preconditions.checkState;
63+
import static com.google.common.collect.ImmutableList.toImmutableList;
6064
import static com.google.common.collect.Sets.intersection;
6165
import static java.lang.String.format;
6266

@@ -271,6 +275,15 @@ protected Void visitQuerySpecification(QuerySpecification querySpec, Context con
271275
return null;
272276
}
273277

278+
@Override
279+
protected Void visitWithQuery(WithQuery withQuery, Context context)
280+
{
281+
context.copyFieldIdsToExploreForWithQuery(withQuery);
282+
process(withQuery.getQuery(), context);
283+
284+
return null;
285+
}
286+
274287
@Override
275288
protected Void visitSampledRelation(SampledRelation sampledRelation, Context context)
276289
{
@@ -493,5 +506,19 @@ private void addFieldIdToExplore(FieldId fieldId)
493506
{
494507
fieldsToExplore.put(fieldId.getRelationId(), fieldId);
495508
}
509+
510+
// Associate the relation from the with clause with the fieldIdsToExplore that we collected for it
511+
// when processing the main part of the query
512+
public void copyFieldIdsToExploreForWithQuery(WithQuery withQuery)
513+
{
514+
QualifiedName name = QualifiedName.of(withQuery.getName().getValue());
515+
List<RelationId> relationIds = fieldsToExplore.keySet().stream()
516+
.filter(key -> key.getSourceNode() instanceof Table && ((Table) key.getSourceNode()).getName().equals(name))
517+
.collect(toImmutableList());
518+
checkState(relationIds.size() == 1, "expected only one matching relationId");
519+
fieldsToExplore.putAll(
520+
RelationId.of(withQuery.getQuery().getQueryBody()),
521+
fieldsToExplore.get(relationIds.get(0)));
522+
}
496523
}
497524
}

presto-main/src/test/java/com/facebook/presto/sql/analyzer/TestUtilizedColumnsAnalyzer.java

+6
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,12 @@ public void testUDF()
519519
ImmutableMap.of(QualifiedObjectName.valueOf("tpch.s1.t1"), ImmutableSet.of("a")));
520520
}
521521

522+
@Test
523+
public void testCteWithExpressionInSelect()
524+
{
525+
assertUtilizedTableColumns("with cte as (select x as c1, z + 1 as c2 from t13) select c1, c2 from (select * from cte)", ImmutableMap.of(QualifiedObjectName.valueOf("tpch.s1.t13"), ImmutableSet.of("x", "z")));
526+
}
527+
522528
private void assertUtilizedTableColumns(@Language("SQL") String query, Map<QualifiedObjectName, Set<String>> expected)
523529
{
524530
transaction(transactionManager, accessControl)

0 commit comments

Comments
 (0)