Skip to content

Commit f69113f

Browse files
joshua-spacetimejoshua-spacetime
authored
Fix filter pushdowns in the presence of rls (#2653)
1 parent 304b488 commit f69113f

2 files changed

Lines changed: 146 additions & 2 deletions

File tree

‎crates/core/src/sql/execute.rs‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -550,6 +550,112 @@ pub(crate) mod tests {
550550
&auth_for_b,
551551
[product![id_for_b]],
552552
);
553+
assert_query_results(
554+
&db,
555+
// Should only return the orders for sender "a"
556+
"select * from users where identity = :sender",
557+
&auth_for_a,
558+
[product![id_for_a]],
559+
);
560+
assert_query_results(
561+
&db,
562+
// Should only return the orders for sender "b"
563+
"select * from users where identity = :sender",
564+
&auth_for_b,
565+
[product![id_for_b]],
566+
);
567+
assert_query_results(
568+
&db,
569+
// Should only return the orders for sender "a"
570+
&format!("select * from users where identity = 0x{}", id_for_a.to_hex()),
571+
&auth_for_a,
572+
[product![id_for_a]],
573+
);
574+
assert_query_results(
575+
&db,
576+
// Should only return the orders for sender "b"
577+
&format!("select * from users where identity = 0x{}", id_for_b.to_hex()),
578+
&auth_for_b,
579+
[product![id_for_b]],
580+
);
581+
assert_query_results(
582+
&db,
583+
// Should only return the orders for sender "a"
584+
&format!(
585+
"select * from users where identity = :sender and identity = 0x{}",
586+
id_for_a.to_hex()
587+
),
588+
&auth_for_a,
589+
[product![id_for_a]],
590+
);
591+
assert_query_results(
592+
&db,
593+
// Should only return the orders for sender "b"
594+
&format!(
595+
"select * from users where identity = :sender and identity = 0x{}",
596+
id_for_b.to_hex()
597+
),
598+
&auth_for_b,
599+
[product![id_for_b]],
600+
);
601+
assert_query_results(
602+
&db,
603+
// Should only return the orders for sender "a"
604+
&format!(
605+
"select * from users where identity = :sender or identity = 0x{}",
606+
id_for_b.to_hex()
607+
),
608+
&auth_for_a,
609+
[product![id_for_a]],
610+
);
611+
assert_query_results(
612+
&db,
613+
// Should only return the orders for sender "b"
614+
&format!(
615+
"select * from users where identity = :sender or identity = 0x{}",
616+
id_for_a.to_hex()
617+
),
618+
&auth_for_b,
619+
[product![id_for_b]],
620+
);
621+
assert_query_results(
622+
&db,
623+
// Should not return any rows.
624+
// Querying as sender "a", but filtering on sender "b".
625+
&format!("select * from users where identity = 0x{}", id_for_b.to_hex()),
626+
&auth_for_a,
627+
[],
628+
);
629+
assert_query_results(
630+
&db,
631+
// Should not return any rows.
632+
// Querying as sender "b", but filtering on sender "a".
633+
&format!("select * from users where identity = 0x{}", id_for_a.to_hex()),
634+
&auth_for_b,
635+
[],
636+
);
637+
assert_query_results(
638+
&db,
639+
// Should not return any rows.
640+
// Querying as sender "a", but filtering on sender "b".
641+
&format!(
642+
"select * from users where identity = :sender and identity = 0x{}",
643+
id_for_b.to_hex()
644+
),
645+
&auth_for_a,
646+
[],
647+
);
648+
assert_query_results(
649+
&db,
650+
// Should not return any rows.
651+
// Querying as sender "b", but filtering on sender "a".
652+
&format!(
653+
"select * from users where identity = :sender and identity = 0x{}",
654+
id_for_a.to_hex()
655+
),
656+
&auth_for_b,
657+
[],
658+
);
553659
assert_query_results(
554660
&db,
555661
// Should only return the orders for sender "a"
@@ -564,6 +670,34 @@ pub(crate) mod tests {
564670
&auth_for_b,
565671
[product![2u64, id_for_b], product![4u64, id_for_b]],
566672
);
673+
assert_query_results(
674+
&db,
675+
// Should only return the orders for sender "a"
676+
"select s.* from users u join sales s on u.identity = s.customer",
677+
&auth_for_a,
678+
[product![1u64, id_for_a], product![3u64, id_for_a]],
679+
);
680+
assert_query_results(
681+
&db,
682+
// Should only return the orders for sender "b"
683+
"select s.* from users u join sales s on u.identity = s.customer",
684+
&auth_for_b,
685+
[product![2u64, id_for_b], product![4u64, id_for_b]],
686+
);
687+
assert_query_results(
688+
&db,
689+
// Should only return the orders for sender "a"
690+
"select s.* from users u join sales s on u.identity = s.customer where u.identity = :sender",
691+
&auth_for_a,
692+
[product![1u64, id_for_a], product![3u64, id_for_a]],
693+
);
694+
assert_query_results(
695+
&db,
696+
// Should only return the orders for sender "b"
697+
"select s.* from users u join sales s on u.identity = s.customer where u.identity = :sender",
698+
&auth_for_b,
699+
[product![2u64, id_for_b], product![4u64, id_for_b]],
700+
);
567701

568702
Ok(())
569703
}

‎crates/physical-plan/src/rules.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -267,9 +267,14 @@ impl RewriteRule for PushConstEq {
267267
type Info = Label;
268268

269269
fn matches(plan: &PhysicalPlan) -> Option<Self::Info> {
270+
// Is this plan a table scan followed by a sequence of filters?
271+
// If so, it's already in a normalized state, so no need to push.
272+
let is_filter = |plan: &PhysicalPlan| {
273+
!plan.any(&|plan| !matches!(plan, PhysicalPlan::TableScan(..) | PhysicalPlan::Filter(..)))
274+
};
270275
if let PhysicalPlan::Filter(input, PhysicalExpr::BinOp(_, expr, value)) = plan {
271276
if let (PhysicalExpr::Field(TupleField { label, .. }), PhysicalExpr::Value(_)) = (&**expr, &**value) {
272-
return (input.has_table_scan(Some(label)) && !input.is_table_scan(None)).then_some(*label);
277+
return (input.has_table_scan(Some(label)) && !is_filter(input)).then_some(*label);
273278
}
274279
}
275280
None
@@ -328,12 +333,17 @@ impl RewriteRule for PushConstAnd {
328333
type Info = Label;
329334

330335
fn matches(plan: &PhysicalPlan) -> Option<Self::Info> {
336+
// Is this plan a table scan followed by a sequence of filters?
337+
// If so, it's already in a normalized state, so no need to push.
338+
let is_filter = |plan: &PhysicalPlan| {
339+
!plan.any(&|plan| !matches!(plan, PhysicalPlan::TableScan(..) | PhysicalPlan::Filter(..)))
340+
};
331341
if let PhysicalPlan::Filter(input, PhysicalExpr::LogOp(LogOp::And, exprs)) = plan {
332342
return exprs.iter().find_map(|expr| {
333343
if let PhysicalExpr::BinOp(_, expr, value) = expr {
334344
if let (PhysicalExpr::Field(TupleField { label, .. }), PhysicalExpr::Value(_)) = (&**expr, &**value)
335345
{
336-
return (input.has_table_scan(Some(label)) && !input.is_table_scan(None)).then_some(*label);
346+
return (input.has_table_scan(Some(label)) && !is_filter(input)).then_some(*label);
337347
}
338348
}
339349
None

0 commit comments

Comments
 (0)