Fix read-only hook test assertions to use as_db_error() instead of to_string()

tokio-postgres 0.7.x Error::to_string() returns generic "db error" without
SQLSTATE details. Use as_db_error().code().code() to properly assert SQLSTATE
25006 in insert/update/delete_blocked_by_read_only_hook tests.

Author: jumpbetweenrows

.githooks/pre-commit

@@ -7,6 +7,9 @@ cargo fmt --check
 echo "→ cargo clippy -p proxy"
 cargo clippy -p proxy -- -D warnings
 
+echo "→ cargo test -p proxy"
+cargo test -p proxy
+
 echo "→ admin-ui typecheck"
 (cd admin-ui && npm run typecheck)
 

proxy/src/engine/mod.rs

@@ -587,6 +587,123 @@ fn select_default_schema(catalog_schemas: &HashMap<String, VirtualCatalogSchema>
         .unwrap_or_else(|| "public".to_string())
 }
 
+// ---------- optimizer rule: fix scan projections for pushed-down filters ----------
+
+/// **DataFusion 52 workaround — revisit on upgrade.**
+///
+/// When PolicyHook injects a `Filter(tenant = 'acme')` above a `TableScan`,
+/// DataFusion's `PushDownFilter` optimizer pushes it into `TableScan.filters`
+/// (because `SqlTable` reports `Exact` pushdown support) and removes the
+/// `Filter` node.  Then `optimize_projections` narrows the scan projection to
+/// `[]` for `COUNT(*)` — but doesn't account for columns needed by the
+/// pushed-down filters in `TableScan.filters`.  The physical planner later
+/// expands the scan projection to include filter columns, creating a schema
+/// mismatch with the logical plan.
+///
+/// This rule runs after all built-in optimizer rules and ensures that any
+/// `TableScan` whose `projection` is `Some(...)` includes all columns
+/// referenced by its pushed-down `filters`.  When extra columns are added, a
+/// wrapping `Projection` strips them so the node's output schema stays
+/// unchanged — preventing mismatches in parent nodes (e.g. `Join`).
+///
+/// **To check on DataFusion upgrade:** run `cargo test --test policy_enforcement
+/// aggregate_with_row_filter` without this rule.  If the test passes, the
+/// workaround can be removed.
+#[derive(Debug)]
+struct ScanFilterProjectionFixRule;
+
+impl datafusion::optimizer::OptimizerRule for ScanFilterProjectionFixRule {
+    fn name(&self) -> &str {
+        "scan_filter_projection_fix"
+    }
+
+    fn apply_order(&self) -> Option<datafusion::optimizer::ApplyOrder> {
+        Some(datafusion::optimizer::ApplyOrder::BottomUp)
+    }
+
+    fn rewrite(
+        &self,
+        plan: datafusion::logical_expr::LogicalPlan,
+        _config: &dyn datafusion::optimizer::OptimizerConfig,
+    ) -> datafusion::error::Result<
+        datafusion::common::tree_node::Transformed<datafusion::logical_expr::LogicalPlan>,
+    > {
+        use datafusion::common::tree_node::Transformed;
+        use datafusion::logical_expr::{Expr, LogicalPlan};
+
+        let LogicalPlan::TableScan(ref scan) = plan else {
+            return Ok(Transformed::no(plan));
+        };
+        // Nothing to fix if no pushed-down filters or projection is already None (all cols).
+        if scan.filters.is_empty() {
+            return Ok(Transformed::no(plan));
+        }
+        let Some(projection) = &scan.projection else {
+            return Ok(Transformed::no(plan));
+        };
+
+        // Collect columns referenced by pushed-down filters that are missing
+        // from the current projection.
+        let source_schema = scan.source.schema();
+        let mut extras: Vec<usize> = Vec::new();
+        for filter_expr in &scan.filters {
+            for col_ref in filter_expr.column_refs() {
+                if let Ok(idx) = source_schema.index_of(&col_ref.name)
+                    && !projection.contains(&idx)
+                    && !extras.contains(&idx)
+                {
+                    extras.push(idx);
+                }
+            }
+        }
+        if extras.is_empty() {
+            return Ok(Transformed::no(plan));
+        }
+
+        // Build expanded scan with filter columns included.
+        let original_proj = projection.clone();
+        let mut new_proj = projection.clone();
+        new_proj.extend(extras);
+        new_proj.sort_unstable();
+
+        let new_scan = datafusion::logical_expr::TableScan::try_new(
+            scan.table_name.clone(),
+            scan.source.clone(),
+            Some(new_proj.clone()),
+            scan.filters.clone(),
+            scan.fetch,
+        )?;
+
+        // Wrap in a Projection that only exposes the original columns so the
+        // output schema is unchanged — parent nodes (Join, Aggregate, etc.)
+        // won't see the extra filter-only columns.
+        let expanded_plan = LogicalPlan::TableScan(new_scan);
+        let expanded_schema = expanded_plan.schema();
+        // Map each original source-column index to its position in the
+        // expanded (sorted) projection so we pull the right DFSchema field.
+        let proj_exprs: Vec<Expr> = original_proj
+            .iter()
+            .map(|&src_idx| {
+                let pos = new_proj
+                    .iter()
+                    .position(|&p| p == src_idx)
+                    .expect("original column must be in expanded projection");
+                let (qualifier, field) = expanded_schema.qualified_field(pos);
+                Expr::Column(datafusion::common::Column::new(
+                    qualifier.cloned(),
+                    field.name(),
+                ))
+            })
+            .collect();
+
+        let projection_plan = LogicalPlan::Projection(
+            datafusion::logical_expr::Projection::try_new(proj_exprs, Arc::new(expanded_plan))?,
+        );
+
+        Ok(Transformed::yes(projection_plan))
+    }
+}
+
 /// Build a SessionContext from local catalog metadata using a shared LazyPool.
 ///
 /// Pool creation is deferred until the first user-table query, so pg_catalog /
@@ -628,8 +745,8 @@ async fn create_session_context_from_catalog(
     let config = SessionConfig::new()
         .with_information_schema(true)
         .with_default_catalog_and_schema("postgres", default_schema);
-
     let mut ctx = SessionContext::new_with_config(config);
+    ctx.add_optimizer_rule(Arc::new(ScanFilterProjectionFixRule));
     ctx.register_catalog("postgres", Arc::new(catalog));
 
     setup_pg_catalog(&ctx, "postgres", ProxyCatalogContext)

proxy/src/hooks/policy.rs

@@ -2,7 +2,6 @@ use arrow_pg::datatypes::df::encode_dataframe;
 use async_trait::async_trait;
 use chrono::Utc;
 use datafusion::common::ScalarValue;
-use datafusion::logical_expr::logical_plan::TableScan;
 use datafusion::logical_expr::registry::FunctionRegistry;
 use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, col, lit};
 use datafusion::prelude::SessionContext;
@@ -995,67 +994,6 @@ impl PolicyEffects {
 
             tracing::debug!(table = %scan.table_name, "PolicyHook: applying row filter");
 
-            // Expand the TableScan projection to include any columns referenced by the
-            // filter expression. When projection=Some([]) (e.g. COUNT(*) optimisation),
-            // the filter's columns would be missing from the scan's output schema without
-            // this expansion, causing a schema mismatch at execution time.
-            let expanded_projection: Option<Vec<usize>> =
-                if let Some(projection) = &scan.projection {
-                    let col_refs = filter_expr.column_refs();
-                    if col_refs.is_empty() {
-                        // lit(false) and similar zero-column-ref filters: no expansion needed.
-                        None
-                    } else {
-                        let full_schema = scan.source.schema();
-                        let table_name_str = scan.table_name.to_string();
-                        let mut extras: Vec<usize> = Vec::new();
-                        for col_ref in &col_refs {
-                            match full_schema.index_of(&col_ref.name) {
-                                Ok(idx) if !projection.contains(&idx) => extras.push(idx),
-                                Ok(_) => {} // already projected
-                                Err(_) => {
-                                    return Err(datafusion::error::DataFusionError::Plan(
-                                        format!(
-                                            "Row filter references column '{}' not found in table '{table_name_str}'",
-                                            col_ref.name
-                                        ),
-                                    ));
-                                }
-                            }
-                        }
-                        if extras.is_empty() {
-                            None
-                        } else {
-                            let mut new_proj = projection.clone();
-                            new_proj.extend(extras);
-                            new_proj.sort_unstable();
-                            Some(new_proj)
-                        }
-                    }
-                } else {
-                    // projection=None means all columns are already available.
-                    None
-                };
-
-            // Rebuild the TableScan with the expanded projection if needed.
-            // (scan borrow ends above; NLL allows moving node here)
-            let node = if let Some(new_projection) = expanded_projection {
-                let LogicalPlan::TableScan(scan_owned) = node else {
-                    unreachable!()
-                };
-                let new_scan = TableScan::try_new(
-                    scan_owned.table_name,
-                    scan_owned.source,
-                    Some(new_projection),
-                    scan_owned.filters,
-                    scan_owned.fetch,
-                )
-                .map_err(|e| datafusion::error::DataFusionError::Plan(e.to_string()))?;
-                LogicalPlan::TableScan(new_scan)
-            } else {
-                node
-            };
-
             let plan_with_filter = LogicalPlanBuilder::from(node)
                 .filter(filter_expr.clone())
                 .and_then(|b| b.build())
@@ -3023,121 +2961,4 @@ mod tests {
             "orders.total remains: {col_names:?}"
         );
     }
-
-    // ---------- apply_row_filters projection expansion ----------
-
-    /// Build a scan plan with an explicit column projection (indices into `columns`).
-    fn build_scan_plan_with_projection(
-        schema_table: &str,
-        columns: Vec<(&str, DataType)>,
-        projection: Option<Vec<usize>>,
-    ) -> LogicalPlan {
-        let fields: Vec<Field> = columns
-            .into_iter()
-            .map(|(name, dt)| Field::new(name, dt, true))
-            .collect();
-        let schema = Arc::new(Schema::new(fields));
-        let table = Arc::new(EmptyTable::new(schema));
-        let source = Arc::new(DefaultTableSource::new(table));
-        LogicalPlanBuilder::scan(schema_table, source, projection)
-            .unwrap()
-            .build()
-            .unwrap()
-    }
-
-    fn make_effects_with_row_filter(
-        key: (&str, &str),
-        filter: datafusion::logical_expr::Expr,
-    ) -> PolicyEffects {
-        let mut effects = PolicyEffects {
-            row_filters: HashMap::new(),
-            column_allow_patterns: HashMap::new(),
-            column_deny_patterns: HashMap::new(),
-            column_masks: HashMap::new(),
-            tables_with_permit: HashSet::new(),
-            denied_by_policy: None,
-        };
-        effects
-            .row_filters
-            .insert((key.0.to_string(), key.1.to_string()), filter);
-        effects
-    }
-
-    #[test]
-    fn test_row_filter_expands_narrow_projection() {
-        // TableScan projecting only col index 0 ("id"), filter references col index 1 ("tenant").
-        // After expansion the projection must contain both 0 and 1.
-        let effects = make_effects_with_row_filter(("s1", "orders"), col("tenant").eq(lit("acme")));
-
-        let plan = build_scan_plan_with_projection(
-            "s1.orders",
-            vec![
-                ("id", DataType::Int32),
-                ("tenant", DataType::Utf8),
-                ("amount", DataType::Int32),
-            ],
-            Some(vec![0]), // only "id" projected
-        );
-
-        let result = effects.apply_row_filters(plan).unwrap();
-
-        // The resulting plan must contain a Filter node above the TableScan.
-        let display = format!("{}", result.display_indent());
-        assert!(
-            display.contains("Filter"),
-            "Expected Filter node in plan: {display}"
-        );
-        // And the projected schema must include "tenant".
-        let schema = result.schema();
-        let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
-        assert!(
-            field_names.contains(&"tenant"),
-            "Expanded projection must include 'tenant': {field_names:?}"
-        );
-    }
-
-    #[test]
-    fn test_row_filter_no_expand_when_all_columns_present() {
-        // projection=None means all columns — no expansion should be attempted.
-        let effects = make_effects_with_row_filter(("s1", "orders"), col("tenant").eq(lit("acme")));
-
-        // Build a scan with projection=None (all columns)
-        let plan = build_scan_plan(
-            "s1.orders",
-            vec![("id", DataType::Int32), ("tenant", DataType::Utf8)],
-        );
-
-        let result = effects.apply_row_filters(plan).unwrap();
-        let display = format!("{}", result.display_indent());
-        assert!(
-            display.contains("Filter"),
-            "Expected Filter node in plan: {display}"
-        );
-    }
-
-    #[test]
-    fn test_row_filter_lit_false_no_expand() {
-        // lit(false) has zero column refs — narrow projection should NOT be expanded.
-        let effects = make_effects_with_row_filter(("s1", "orders"), lit(false));
-
-        let plan = build_scan_plan_with_projection(
-            "s1.orders",
-            vec![("id", DataType::Int32), ("tenant", DataType::Utf8)],
-            Some(vec![0]), // only "id" projected — must NOT expand for lit(false)
-        );
-
-        let result = effects.apply_row_filters(plan).unwrap();
-        let display = format!("{}", result.display_indent());
-        assert!(
-            display.contains("Filter"),
-            "Expected Filter node in plan: {display}"
-        );
-        // "tenant" must NOT appear in the projected schema (no expansion for lit(false))
-        let schema = result.schema();
-        let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
-        assert!(
-            !field_names.contains(&"tenant"),
-            "lit(false) filter must not expand projection: {field_names:?}"
-        );
-    }
 }