diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 531c4fb26..d7e80a11b 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -947,6 +947,14 @@ impl PaimonTableScan { self.limit } + /// Planned under the server's rules: counts and bounds precede them. + fn has_query_auth_rules(&self) -> bool { + self.planned_partitions + .iter() + .flat_map(|splits| splits.iter()) + .any(DataSplit::has_query_auth_rules) + } + fn manifest_column_statistics(&self, partitions: &[Arc<[DataSplit]>]) -> Vec { if self.read_type.len() != self.schema().fields().len() { return Statistics::unknown_column(&self.schema()); @@ -1182,6 +1190,10 @@ impl ExecutionPlan for PaimonTableScan { } fn partition_statistics(&self, partition: Option) -> DFResult> { + // They include the rows the server's rules drop. + if self.has_query_auth_rules() { + return Ok(Arc::new(Statistics::new_unknown(&self.schema()))); + } let partitions: &[Arc<[DataSplit]>] = match partition { Some(idx) => std::slice::from_ref(&self.planned_partitions[idx]), None => &self.planned_partitions, @@ -1229,6 +1241,11 @@ impl DisplayAs for PaimonTableScan { f: &mut std::fmt::Formatter, ) -> std::fmt::Result { write!(f, "PaimonTableScan: table={}", self.table.identifier())?; + let restricted = self.has_query_auth_rules(); + if restricted { + // Split and file counts describe the data before the server's rules. + write!(f, ", query-auth=restricted")?; + } let total_splits: usize = self.planned_partitions.iter().map(|p| p.len()).sum(); let total_files: usize = self @@ -1237,11 +1254,13 @@ impl DisplayAs for PaimonTableScan { .flat_map(|p| p.iter()) .map(|s| s.data_files().len()) .sum(); - write!( - f, - ", partitions={}, splits={total_splits}, files={total_files}", - self.planned_partitions.len() - )?; + if !restricted { + write!( + f, + ", partitions={}, splits={total_splits}, files={total_files}", + self.planned_partitions.len() + )?; + } let columns = self .read_type @@ -1255,7 +1274,7 @@ impl DisplayAs for PaimonTableScan { if let Some(limit) = self.limit { write!(f, ", limit={limit}")?; } - if let Some(ref trace) = self.scan_trace { + if let Some(trace) = self.scan_trace.as_ref().filter(|_| !restricted) { write!(f, ", trace={trace}")?; } if let Some(ref pushed_variants) = self.pushed_variants { diff --git a/crates/integrations/datafusion/tests/query_auth.rs b/crates/integrations/datafusion/tests/query_auth.rs new file mode 100644 index 000000000..e9154c75c --- /dev/null +++ b/crates/integrations/datafusion/tests/query_auth.rs @@ -0,0 +1,227 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! SQL over a `query-auth.enabled` table whose REST server sets a row filter +//! for the current user. + +mod common; + +#[path = "../../../paimon/tests/mock_server.rs"] +mod mock_server; + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::{Array, Int32Array, Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as ArrowDataType, Field, Schema as ArrowSchema}; +use paimon::api::{AuthTableQueryResponse, ConfigResponse}; +use paimon::catalog::{Catalog, Identifier, RESTCatalog}; +use paimon::spec::{DataType, IntType, Schema, VarCharType}; +use paimon::{CatalogOptions, FileSystemCatalog, Options}; +use paimon_datafusion::SQLContext; +use serde_json::json; + +use mock_server::{start_mock_server, RESTServer}; + +const TABLE: &str = "paimon.default.people"; + +fn schema(options: &[(&str, &str)]) -> Schema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("secret", DataType::VarChar(VarCharType::new(255).unwrap())); + for (key, value) in options { + builder = builder.option(*key, *value); + } + builder.build().unwrap() +} + +/// Ids 1..=10 written through a filesystem catalog, served by a mock REST +/// catalog that admits `id > 6` and applies `column_masking`. +async fn restricted_people( + column_masking: Option>, +) -> (tempfile::TempDir, RESTServer, SQLContext) { + let tmp = tempfile::tempdir().unwrap(); + let mut fs_options = Options::new(); + fs_options.set( + CatalogOptions::WAREHOUSE, + format!("file://{}", tmp.path().display()), + ); + let fs_catalog = FileSystemCatalog::new(fs_options).unwrap(); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let identifier = Identifier::new("default", "people"); + fs_catalog + .create_table(&identifier, schema(&[]), false) + .await + .unwrap(); + let table = fs_catalog.get_table(&identifier).await.unwrap(); + let names = [ + "alice", "bob", "carol", "dave", "erin", "frank", "grace", "heidi", "ivan", "judy", + ]; + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowDataType::Int32, true), + Field::new("name", ArrowDataType::Utf8, true), + Field::new("secret", ArrowDataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from_iter_values(1..=10)), + Arc::new(StringArray::from_iter_values(names)), + Arc::new(StringArray::from_iter_values( + names.iter().map(|n| format!("{n}-ssn")), + )), + ], + ) + .unwrap(); + let write_builder = table.new_write_builder(); + let mut write = write_builder.new_write().unwrap(); + write.write_arrow_batch(&batch).await.unwrap(); + let messages = write.prepare_commit().await.unwrap(); + write_builder.new_commit().commit(messages).await.unwrap(); + + let server = start_mock_server( + "test_warehouse".to_string(), + tmp.path().to_string_lossy().into_owned(), + ConfigResponse::new(HashMap::from([( + CatalogOptions::PREFIX.to_string(), + "mock-test".to_string(), + )])), + vec!["default".to_string()], + ) + .await; + server.add_table_with_schema( + "default", + "people", + schema(&[("query-auth.enabled", "true")]), + table.location(), + ); + server.set_auth_response( + "default", + "people", + AuthTableQueryResponse { + filter: Some(vec![json!({ + "kind": "LEAF", + "transform": { + "name": "FIELD_REF", + "fieldRef": {"index": 0, "name": "id", "type": "INT"}, + }, + "function": "GREATER_THAN", + "literals": [6], + }) + .to_string()]), + column_masking, + }, + ); + + let mut options = Options::new(); + options.set(CatalogOptions::URI, server.url().unwrap()); + options.set(CatalogOptions::WAREHOUSE, "test_warehouse"); + options.set(CatalogOptions::TOKEN_PROVIDER, "bear"); + options.set(CatalogOptions::TOKEN, "test-token"); + let catalog = Arc::new(RESTCatalog::new(options, true).await.unwrap()); + let mut context = SQLContext::new(); + context.register_catalog("paimon", catalog).await.unwrap(); + (tmp, server, context) +} + +async fn query(context: &SQLContext, sql: &str) -> Vec { + context.sql(sql).await.unwrap().collect().await.unwrap() +} + +fn id_names(batches: &[RecordBatch]) -> Vec<(i32, String)> { + let mut rows = Vec::new(); + for batch in batches { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..batch.num_rows() { + rows.push(( + ids.value(row), + common::string_value(batch.column(1).as_ref(), row).to_string(), + )); + } + } + rows.sort(); + rows +} + +/// The one value of a one-row, one-column result. +fn single(batches: &[RecordBatch]) -> Arc { + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 1, + "{batches:?}" + ); + let batch = batches.iter().find(|b| b.num_rows() == 1).unwrap(); + assert_eq!(batch.num_columns(), 1); + Arc::clone(batch.column(0)) +} + +fn count(batches: &[RecordBatch]) -> i64 { + single(batches) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_query_auth_queries_see_only_the_rows_the_filter_admits() { + let (_tmp, _server, context) = restricted_people(None).await; + + assert_eq!( + id_names(&query(&context, &format!("SELECT id, name FROM {TABLE}")).await), + vec![ + (7, "grace".to_string()), + (8, "heidi".to_string()), + (9, "ivan".to_string()), + (10, "judy".to_string()), + ] + ); + // Manifest statistics would answer 10 and 1. + assert_eq!( + count(&query(&context, &format!("SELECT COUNT(*) FROM {TABLE}")).await), + 4 + ); + let min = single(&query(&context, &format!("SELECT MIN(id) FROM {TABLE}")).await); + assert_eq!( + min.as_any().downcast_ref::().unwrap().value(0), + 7 + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_query_auth_explain_shows_the_restriction_but_not_the_files() { + let (_tmp, _server, context) = restricted_people(None).await; + + let plan = datafusion::arrow::util::pretty::pretty_format_batches( + &query( + &context, + &format!("EXPLAIN SELECT id, name FROM {TABLE} WHERE name = 'grace'"), + ) + .await, + ) + .unwrap() + .to_string(); + assert!(plan.contains("query-auth=restricted"), "{plan}"); + assert!(!plan.contains("files="), "{plan}"); +} diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index d9e958b6a..5664e2bba 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -417,10 +417,12 @@ impl Table { } /// Whether this user may read this table; `None` when it is not - /// `query-auth.enabled`. `query_auth` is the loaded option, as in Java. + /// `query-auth.enabled`. `query_auth` is the loaded option, as in Java; + /// `select` is what the read touches, `None` for the whole table. pub(crate) async fn authorize_read( &self, query_auth: bool, + select: Option>, ) -> Result>> { let local = CoreOptions::new(self.schema.options()); let Some(rest_env) = &self.rest_env else { @@ -451,13 +453,12 @@ impl Table { query_auth::unsupported("this table handle was assembled rather than loaded") })?; - // Naming a system column here would fail the server's column check. let response = rest_env - .table_query_auth(self.schema.id(), self.schema.fields(), None) + .table_query_auth(self.schema.id(), self.schema.fields(), select.clone()) .await?; - Ok(Some(std::sync::Arc::new(query_auth::QueryAuthGrant::new( - response, session, - )))) + Ok(Some(std::sync::Arc::new( + query_auth::QueryAuthGrant::parse(&response, session, select, self.schema.fields())?, + ))) } /// Handed out once per catalog-loaded table; wraps only after 2^64 loads. diff --git a/crates/paimon/src/table/query_auth.rs b/crates/paimon/src/table/query_auth.rs index fb24d3bae..e767ed205 100644 --- a/crates/paimon/src/table/query_auth.rs +++ b/crates/paimon/src/table/query_auth.rs @@ -17,24 +17,51 @@ //! What the REST server authorized a user to read from one table. +mod rules; + use crate::api::AuthTableQueryResponse; +pub(crate) use rules::{filter_batch, Rules}; -/// The server's answer for one user on one table, kept unparsed; `session` -/// ties it to the handle that asked, as the response names no table or user. +/// The server's answer for one user on one table; `session` ties it to the +/// handle that asked, as the response names no table or user. #[derive(Debug, PartialEq)] pub(crate) struct QueryAuthGrant { - response: AuthTableQueryResponse, session: u64, + /// The columns the request asked about; `None` asked about the whole table. + select: Option>, + rules: Rules, } impl QueryAuthGrant { - pub(crate) fn new(response: AuthTableQueryResponse, session: u64) -> Self { - Self { response, session } + /// Parses the rules against `fields`, the schema the server ruled on. + pub(crate) fn parse( + response: &AuthTableQueryResponse, + session: u64, + select: Option>, + fields: &[crate::spec::DataField], + ) -> crate::Result { + Ok(Self::new(session, select, Rules::parse(response, fields)?)) } - /// The only case this client can serve. + pub(crate) fn new(session: u64, select: Option>, rules: Rules) -> Self { + Self { + session, + select, + rules, + } + } + + /// No row filter and no masking. pub(crate) fn is_unrestricted(&self) -> bool { - self.response.is_unrestricted() + self.rules.is_empty() + } + + pub(crate) fn rules(&self) -> &Rules { + &self.rules + } + + pub(crate) fn select(&self) -> Option<&[String]> { + self.select.as_deref() } /// A view of another schema is not the one the server ruled on. @@ -115,6 +142,28 @@ fn contains(wide: &crate::spec::DataType, narrow: &crate::spec::DataType) -> boo } } +/// Every leaf's column name, system columns included. +pub(crate) fn leaf_names( + predicates: &[crate::spec::Predicate], +) -> std::collections::HashSet { + fn collect(predicate: &crate::spec::Predicate, out: &mut std::collections::HashSet) { + use crate::spec::Predicate; + match predicate { + Predicate::Leaf { column, .. } => { + out.insert(column.clone()); + } + Predicate::And(children) | Predicate::Or(children) => { + children.iter().for_each(|child| collect(child, out)); + } + Predicate::Not(inner) => collect(inner, out), + Predicate::AlwaysTrue | Predicate::AlwaysFalse => {} + } + } + let mut out = std::collections::HashSet::new(); + predicates.iter().for_each(|p| collect(p, &mut out)); + out +} + /// A refusal naming the option, so callers never match on prose. pub(crate) fn unsupported(reason: &str) -> crate::Error { crate::Error::Unsupported { @@ -177,10 +226,8 @@ mod tests { async fn test_a_grant_is_pinned_to_the_handle_that_obtained_it() { let a = crate::table::rest_query_auth_table().await; let b = crate::table::rest_query_auth_table().await; - let grant = super::QueryAuthGrant::new( - crate::api::AuthTableQueryResponse::default(), - a.query_auth_session().unwrap(), - ); + let grant = + super::QueryAuthGrant::new(a.query_auth_session().unwrap(), None, Default::default()); assert!(grant.matches_table(&a)); assert!( !grant.matches_table(&b), @@ -206,7 +253,7 @@ mod tests { "1".to_string(), )])); assert!(!table.is_time_traveled(), "{selector} sets no flag"); - let err = table.authorize_read(true).await.unwrap_err(); + let err = table.authorize_read(true, None).await.unwrap_err(); assert!( matches!(err, crate::Error::Unsupported { ref message } if message.contains("time-travelled or branch read")), @@ -240,7 +287,7 @@ mod tests { ("scan.snapshot-id".to_string(), "invalid".to_string()), ])); assert!(table.reads_another_schema().unwrap()); - assert!(table.authorize_read(false).await.unwrap().is_none()); + assert!(table.authorize_read(false, None).await.unwrap().is_none()); } #[tokio::test] @@ -257,8 +304,9 @@ mod tests { async fn test_a_grant_does_not_cross_into_a_travelled_or_branch_view() { let table = crate::table::rest_query_auth_table().await; let grant = super::QueryAuthGrant::new( - crate::api::AuthTableQueryResponse::default(), table.query_auth_session().unwrap(), + None, + Default::default(), ); assert!(grant.matches_table(&table)); @@ -497,7 +545,7 @@ mod tests { async fn test_time_travelled_or_branch_read_is_refused() { let mut travelled = rest_query_auth_table().await; travelled.time_traveled = true; - let err = travelled.authorize_read(true).await.unwrap_err(); + let err = travelled.authorize_read(true, None).await.unwrap_err(); assert!( matches!(err, crate::Error::Unsupported { ref message } if message.contains("time-travelled or branch read")), @@ -506,6 +554,6 @@ mod tests { let mut branch = rest_query_auth_table().await; branch.branch_reference = true; - assert!(branch.authorize_read(true).await.is_err()); + assert!(branch.authorize_read(true, None).await.is_err()); } } diff --git a/crates/paimon/src/table/query_auth/rules.rs b/crates/paimon/src/table/query_auth/rules.rs new file mode 100644 index 000000000..207103903 --- /dev/null +++ b/crates/paimon/src/table/query_auth/rules.rs @@ -0,0 +1,699 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The server's row filter, parsed against the schema it ruled on and applied +//! to read batches, as Java `TableQueryAuthResult` does. + +use super::unsupported; +use crate::api::AuthTableQueryResponse; +use crate::arrow::residual::{evaluate_exact_leaf_predicate, sanitize_filter_mask}; +use crate::spec::{DataField, Predicate}; +use crate::Result; +use arrow_array::{Array, ArrayRef, BooleanArray, Float32Array, Float64Array, RecordBatch}; +use std::collections::HashSet; +use std::sync::Arc; + +/// The rules of one grant; leaf indices are table-schema positions. +#[derive(Debug, Default, PartialEq)] +pub(crate) struct Rules { + pub(crate) filters: Vec, +} + +impl Rules { + /// Validated like Java `validateAgainstSchema`: anything unknown fails closed. + pub(crate) fn parse(response: &AuthTableQueryResponse, fields: &[DataField]) -> Result { + let mut filters = Vec::new(); + for json in response.filter.iter().flatten() { + if json.is_empty() { + return Err(unsupported("the server sent an empty row filter")); + } + let filter = + Predicate::from_rest_json(json, fields).map_err( + |e| match mentioned_system_column(json) { + Some(name) => unsupported(&format!( + "the server's row filter reads the system column '{name}', which a \ + query-auth read never projects" + )), + None => unsupported(&format!("cannot parse the server's row filter: {e}")), + }, + )?; + filters.push(filter); + } + + // Not applied yet: refusing beats returning raw values. + if response + .column_masking + .as_ref() + .is_some_and(|masks| !masks.is_empty()) + { + return Err(unsupported("this client does not apply column masking yet")); + } + Ok(Self { filters }) + } + + pub(crate) fn is_empty(&self) -> bool { + self.filters.is_empty() + } + + /// Table-schema indices the row filter reads. + pub(crate) fn filter_columns(&self) -> HashSet { + let mut out = HashSet::new(); + for filter in &self.filters { + filter.collect_leaf_field_indices(&mut out); + } + out + } +} + +fn mentioned_system_column(json: &str) -> Option<&'static str> { + [ + crate::spec::ROW_ID_FIELD_NAME, + crate::spec::SEQUENCE_NUMBER_FIELD_NAME, + crate::spec::VALUE_KIND_FIELD_NAME, + crate::spec::ROW_KIND_FIELD_NAME, + ] + .into_iter() + .find(|name| json.contains(name)) +} + +// --------------------------------------------------------------------------- +// Row filter +// --------------------------------------------------------------------------- + +/// Keeps the rows passing every filter. `batch` columns pair 1:1 with +/// `batch_fields`; leaf indices point into `schema_fields`. Anything that +/// cannot be evaluated is an error, never a kept row. +pub(crate) fn filter_batch( + batch: &RecordBatch, + filters: &[Predicate], + schema_fields: &[DataField], + batch_fields: &[DataField], +) -> Result { + let mut keep: Option = None; + for filter in filters { + let mask = rule_mask(batch, filter, schema_fields, batch_fields)?; + keep = Some(match keep { + Some(existing) => combine(&existing, &mask, false)?, + None => mask, + }); + } + let Some(keep) = keep else { + return Ok(batch.clone()); + }; + arrow_select::filter::filter_record_batch(batch, &keep).map_err(|e| eval_err(&e)) +} + +/// Java leaf semantics: a null value fails every leaf but `IS NULL`, for the +/// negated functions too. The REST parser builds no `NOT`, so two-valued +/// AND/OR over those leaves is exactly Java's `test`. +fn rule_mask( + batch: &RecordBatch, + predicate: &Predicate, + schema_fields: &[DataField], + batch_fields: &[DataField], +) -> Result { + match predicate { + Predicate::AlwaysTrue => Ok(BooleanArray::from(vec![true; batch.num_rows()])), + Predicate::AlwaysFalse => Ok(BooleanArray::from(vec![false; batch.num_rows()])), + Predicate::And(children) | Predicate::Or(children) => { + let or = matches!(predicate, Predicate::Or(_)); + let mut combined: Option = None; + for child in children { + let mask = rule_mask(batch, child, schema_fields, batch_fields)?; + combined = Some(match combined { + Some(existing) => combine(&existing, &mask, or)?, + None => mask, + }); + } + combined.ok_or_else(|| unsupported("the server's row filter has an empty AND/OR")) + } + // Its null semantics differ from Java's negated leaves. + Predicate::Not(_) => Err(unsupported("the server's row filter has a NOT")), + Predicate::Leaf { + index, + op, + literals, + .. + } => { + let field = schema_fields + .get(*index) + .ok_or_else(|| unsupported("the server's row filter reads an unknown column"))?; + let column = column_of(batch, field, batch_fields)?; + let column = canonicalize_nan(&column); + let mask = evaluate_exact_leaf_predicate(&column, field.data_type(), *op, literals) + .map_err(|e| eval_err(&e))?; + Ok(sanitize_filter_mask(mask)) + } + } +} + +fn combine(left: &BooleanArray, right: &BooleanArray, or: bool) -> Result { + let combined = if or { + arrow_arith::boolean::or(left, right) + } else { + arrow_arith::boolean::and(left, right) + }; + combined.map_err(|e| eval_err(&e)) +} + +/// The column holding `field`, matched by id and name so an alias cannot stand in. +fn column_of( + batch: &RecordBatch, + field: &DataField, + batch_fields: &[DataField], +) -> Result { + let position = batch_fields + .iter() + .position(|f| f.id() == field.id() && f.name() == field.name()) + .ok_or_else(|| { + unsupported(&format!( + "the read does not carry '{}', which the server's rules need", + field.name() + )) + })?; + let column = batch.column(position); + let declared = crate::arrow::paimon_type_to_arrow(field.data_type())?; + if column.data_type() == &declared { + return Ok(Arc::clone(column)); + } + arrow_cast::cast(column, &declared).map_err(|e| eval_err(&e)) +} + +/// Java's `Float.compare` puts every NaN above every number; Arrow's total +/// order puts a negative NaN below them, so `f < 0` would admit it. +fn canonicalize_nan(column: &ArrayRef) -> ArrayRef { + match column.data_type() { + arrow_schema::DataType::Float32 => match column.as_any().downcast_ref::() { + Some(values) if values.iter().any(|v| v.is_some_and(f32::is_nan)) => { + Arc::new(values.unary::<_, arrow_array::types::Float32Type>(|v| { + if v.is_nan() { + f32::NAN + } else { + v + } + })) + } + _ => Arc::clone(column), + }, + arrow_schema::DataType::Float64 => match column.as_any().downcast_ref::() { + Some(values) if values.iter().any(|v| v.is_some_and(f64::is_nan)) => { + Arc::new(values.unary::<_, arrow_array::types::Float64Type>(|v| { + if v.is_nan() { + f64::NAN + } else { + v + } + })) + } + _ => Arc::clone(column), + }, + _ => Arc::clone(column), + } +} + +fn eval_err(e: &dyn std::fmt::Display) -> crate::Error { + unsupported(&format!("cannot evaluate the server's rules: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{ + BigIntType, BooleanType, DataType, DateType, Datum, DecimalType, DoubleType, FloatType, + IntType, PredicateBuilder, PredicateOperator, TimestampType, VarBinaryType, VarCharType, + }; + use arrow_array::{ + Date32Array, Decimal128Array, Int32Array, Int64Array, RecordBatchOptions, StringArray, + TimestampMillisecondArray, + }; + use arrow_schema::{Field, Schema}; + + const NULL: &str = r#"{"name":"NULL"}"#; + + fn field(id: i32, name: &str, data_type: DataType) -> DataField { + DataField::new(id, name.to_string(), data_type) + } + + fn string() -> DataType { + DataType::VarChar(VarCharType::string_type()) + } + + fn batch(fields: &[DataField], columns: Vec) -> RecordBatch { + let schema: Vec = fields + .iter() + .map(|f| { + let data_type = crate::arrow::paimon_type_to_arrow(f.data_type()).unwrap(); + Field::new(f.name(), data_type, f.data_type().is_nullable()) + }) + .collect(); + RecordBatch::try_new(Arc::new(Schema::new(schema)), columns).unwrap() + } + + fn parse(filters: &[&str], masks: &[(&str, &str)], fields: &[DataField]) -> Result { + let response = AuthTableQueryResponse { + filter: Some(filters.iter().map(|f| f.to_string()).collect()), + column_masking: Some( + masks + .iter() + .map(|(column, json)| (column.to_string(), json.to_string())) + .collect(), + ), + }; + Rules::parse(&response, fields) + } + + fn refusal(result: Result) -> String { + match result { + Err(crate::Error::Unsupported { message }) => message, + other => panic!("expected a refusal, got {other:?}"), + } + } + + /// Java `LeafPredicate` JSON; the field resolves by name and the local type wins. + fn leaf(function: &str, field: &str, literals: &str) -> String { + format!( + r#"{{"kind":"LEAF","transform":{{"name":"FIELD_REF","fieldRef":{{"index":0,"name":"{field}","type":"INT"}}}},"function":"{function}","literals":{literals}}}"# + ) + } + + fn compound(function: &str, children: &[String]) -> String { + format!( + r#"{{"kind":"COMPOUND","function":"{function}","children":[{}]}}"#, + children.join(",") + ) + } + + fn constant(function: &str) -> String { + format!( + r#"{{"kind":"LEAF","transform":{{"name":"NULL"}},"function":"{function}","literals":[]}}"# + ) + } + + fn ids(batch: &RecordBatch) -> Vec { + let ids = batch.column(0).as_any().downcast_ref::(); + ids.unwrap().values().to_vec() + } + + // --------------------------------------------------------------------- + // Parsing + // --------------------------------------------------------------------- + + fn rule_fields() -> Vec { + let bytes = VarBinaryType::new(VarBinaryType::MAX_LENGTH).unwrap(); + let not_null = VarCharType::with_nullable(false, VarCharType::MAX_LENGTH).unwrap(); + vec![ + field(0, "id", DataType::Int(IntType::new())), + field(1, "name", string()), + field(2, "alias", string()), + field(3, "bin", DataType::VarBinary(bytes)), + field(4, "nn", DataType::VarChar(not_null)), + ] + } + + #[test] + fn test_parse_reads_the_rules() { + let fields = rule_fields(); + let rules = parse( + &[ + &leaf("GREATER_THAN", "id", "[1]"), + &leaf("IS_NOT_NULL", "alias", "[]"), + ], + &[], + &fields, + ) + .unwrap(); + assert_eq!(rules.filters.len(), 2); + assert_eq!(rules.filter_columns(), HashSet::from([0, 2])); + + let none = AuthTableQueryResponse::default(); + assert!(Rules::parse(&none, &fields).unwrap().is_empty()); + assert!(parse(&[], &[], &fields).unwrap().is_empty()); + } + + #[test] + fn test_any_mask_is_refused() { + let message = refusal(parse(&[], &[("name", NULL)], &rule_fields())); + assert!(message.contains("column masking"), "{message}"); + } + + #[test] + fn test_parse_fails_closed() { + let fields = rule_fields(); + for (filter, reason) in [ + (String::new(), "empty row filter"), + (" ".to_string(), "cannot parse"), + ("null".to_string(), "cannot parse"), + ("{".to_string(), "cannot parse"), + (leaf("EQUAL", "missing", "[1]"), "unknown field `missing`"), + (leaf("EQUAL", "id", r#"["one"]"#), "cannot parse"), + (compound("AND", &[]), "empty children"), + ] { + let message = refusal(parse(&[&filter], &[], &fields)); + assert!(message.contains(reason), "{filter:?}: {message}"); + } + // One bad entry spoils the list. + let good = leaf("EQUAL", "id", "[1]"); + assert!(refusal(parse(&[&good, ""], &[], &fields)).contains("empty row filter")); + } + + #[test] + fn test_a_filter_on_a_system_column_is_refused() { + let fields = rule_fields(); + for column in ["_ROW_ID", "_SEQUENCE_NUMBER", "_VALUE_KIND", "rowkind"] { + let message = refusal(parse(&[&leaf("EQUAL", column, "[1]")], &[], &fields)); + assert!( + message.contains(&format!("system column '{column}'")), + "{message}" + ); + } + } + + // --------------------------------------------------------------------- + // Row filter + // --------------------------------------------------------------------- + + /// `k` names the row; every other column is null in the last one. + fn typed_fields() -> Vec { + vec![ + field(0, "k", DataType::Int(IntType::with_nullable(false))), + field(1, "i", DataType::Int(IntType::new())), + field(2, "b", DataType::BigInt(BigIntType::new())), + field(3, "d", DataType::Double(DoubleType::new())), + field(4, "f", DataType::Float(FloatType::new())), + field(5, "s", string()), + field(6, "m", DataType::Decimal(DecimalType::new(5, 2).unwrap())), + field(7, "dt", DataType::Date(DateType::new())), + field(8, "ts", DataType::Timestamp(TimestampType::new(3).unwrap())), + field(9, "flag", DataType::Boolean(BooleanType::new())), + ] + } + + /// Four values, then a null. + fn with_null(values: [T; 4]) -> Vec> { + values.into_iter().map(Some).chain([None]).collect() + } + + fn typed_batch() -> RecordBatch { + let decimals = Decimal128Array::from(with_null([110, 200, 250, 1000])); + batch( + &typed_fields(), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4])), + Arc::new(Int32Array::from(with_null([1, 2, 3, 4]))), + Arc::new(Int64Array::from(with_null([10, 20, 30, 40]))), + Arc::new(Float64Array::from(with_null([-1.5, 0.0, -0.0, f64::NAN]))), + Arc::new(Float32Array::from(with_null([-1.5, 0.0, -0.0, f32::NAN]))), + Arc::new(StringArray::from(with_null([ + "apple", "apricot", "banana", "a_c", + ]))), + Arc::new(decimals.with_precision_and_scale(5, 2).unwrap()), + Arc::new(Date32Array::from(with_null([19000, 19001, 19002, 19003]))), + Arc::new(TimestampMillisecondArray::from(with_null([ + 1000, 2000, 3000, 4000, + ]))), + Arc::new(BooleanArray::from(with_null([true, false, true, false]))), + ], + ) + } + + /// The rows of the typed batch that `filters` keep. + fn kept(filters: &[Predicate]) -> Vec { + let fields = typed_fields(); + ids(&filter_batch(&typed_batch(), filters, &fields, &fields).unwrap()) + } + + #[test] + fn test_filter_matrix_follows_java_leaf_semantics() { + let fields = typed_fields(); + // A null value fails every leaf but IS_NULL, the negated ones too. + let cases = [ + (leaf("EQUAL", "i", "[2]"), vec![1]), + (leaf("NOT_EQUAL", "i", "[2]"), vec![0, 2, 3]), + (leaf("LESS_THAN", "i", "[3]"), vec![0, 1]), + (leaf("GREATER_OR_EQUAL", "i", "[3]"), vec![2, 3]), + (leaf("IN", "i", "[1,3,null]"), vec![0, 2]), + (leaf("NOT_IN", "i", "[1,3]"), vec![1, 3]), + (leaf("NOT_IN", "i", "[1,null]"), vec![]), + (leaf("BETWEEN", "i", "[2,3]"), vec![1, 2]), + (leaf("NOT_BETWEEN", "i", "[2,3]"), vec![0, 3]), + (leaf("IS_NULL", "i", "[]"), vec![4]), + (leaf("IS_NOT_NULL", "i", "[]"), vec![0, 1, 2, 3]), + (leaf("EQUAL", "i", "[null]"), vec![]), + (leaf("GREATER_OR_EQUAL", "b", "[30]"), vec![2, 3]), + (leaf("NOT_IN", "b", "[10,40]"), vec![1, 2]), + (leaf("BETWEEN", "b", "[15,35]"), vec![1, 2]), + // `Double.compare`: -0.0 sorts below 0.0, NaN above every number. + (leaf("EQUAL", "d", "[0.0]"), vec![1]), + (leaf("EQUAL", "d", "[-0.0]"), vec![2]), + (leaf("NOT_EQUAL", "d", "[0.0]"), vec![0, 2, 3]), + (leaf("LESS_THAN", "d", "[0.0]"), vec![0, 2]), + (leaf("GREATER_OR_EQUAL", "d", "[0.0]"), vec![1, 3]), + (leaf("IN", "d", "[0.0,-1.5]"), vec![0, 1]), + (leaf("NOT_IN", "d", "[0.0]"), vec![0, 2, 3]), + (leaf("BETWEEN", "d", "[-1.0,1.0]"), vec![1, 2]), + (leaf("NOT_BETWEEN", "d", "[-1.0,1.0]"), vec![0, 3]), + (leaf("EQUAL", "f", "[-1.5]"), vec![0]), + (leaf("NOT_EQUAL", "f", "[-0.0]"), vec![0, 1, 3]), + (leaf("LESS_THAN", "f", "[0.0]"), vec![0, 2]), + (leaf("GREATER_OR_EQUAL", "f", "[0.0]"), vec![1, 3]), + (leaf("EQUAL", "s", r#"["banana"]"#), vec![2]), + (leaf("NOT_EQUAL", "s", r#"["banana"]"#), vec![0, 1, 3]), + (leaf("LESS_THAN", "s", r#"["apricot"]"#), vec![0, 3]), + (leaf("GREATER_OR_EQUAL", "s", r#"["apricot"]"#), vec![1, 2]), + (leaf("IN", "s", r#"["apple","banana"]"#), vec![0, 2]), + (leaf("NOT_IN", "s", r#"["apple"]"#), vec![1, 2, 3]), + (leaf("BETWEEN", "s", r#"["apple","b"]"#), vec![0, 1]), + (leaf("NOT_BETWEEN", "s", r#"["apple","b"]"#), vec![2, 3]), + (leaf("STARTS_WITH", "s", r#"["ap"]"#), vec![0, 1]), + (leaf("LIKE", "s", r#"["%an%"]"#), vec![2]), + (leaf("LIKE", "s", r#"["a_p%"]"#), vec![0]), + (leaf("LIKE", "s", r#"["a\\_c"]"#), vec![3]), + (leaf("IS_NULL", "s", "[]"), vec![4]), + (leaf("EQUAL", "flag", "[true]"), vec![0, 2]), + (leaf("NOT_EQUAL", "flag", "[true]"), vec![1, 3]), + (leaf("IN", "flag", "[false]"), vec![1, 3]), + (leaf("IS_NOT_NULL", "flag", "[]"), vec![0, 1, 2, 3]), + ]; + for (json, expected) in cases { + let rules = parse(&[&json], &[], &fields).unwrap(); + assert_eq!(kept(&rules.filters), expected, "{json}"); + } + } + + #[test] + fn test_filter_compares_decimals_dates_and_timestamps_by_value() { + // The REST parser takes no such literal, so these leaves are built directly. + let fields = typed_fields(); + let b = PredicateBuilder::new(&fields); + let dec = |unscaled, scale| Datum::Decimal { + unscaled, + precision: 10, + scale, + }; + let ts = |millis, nanos| Datum::Timestamp { millis, nanos }; + let cases = [ + (b.equal("m", dec(2, 0)), vec![1]), + (b.equal("m", dec(20, 1)), vec![1]), + (b.not_equal("m", dec(25, 1)), vec![0, 1, 3]), + (b.greater_than("m", dec(2001, 3)), vec![2, 3]), + (b.less_or_equal("m", dec(11, 1)), vec![0]), + (b.between("m", dec(11, 1), dec(25, 1)), vec![0, 1, 2]), + (b.not_between("m", dec(11, 1), dec(25, 1)), vec![3]), + (b.is_in("m", vec![dec(10, 0), dec(1100, 3)]), vec![0, 3]), + (b.is_not_in("m", vec![dec(2, 0)]), vec![0, 2, 3]), + (b.greater_than("dt", Datum::Date(19001)), vec![2, 3]), + ( + b.between("dt", Datum::Date(19000), Datum::Date(19001)), + vec![0, 1], + ), + (b.not_equal("dt", Datum::Date(19000)), vec![1, 2, 3]), + ( + b.is_in("dt", vec![Datum::Date(19000), Datum::Date(19003)]), + vec![0, 3], + ), + (b.is_null("dt"), vec![4]), + (b.less_than("ts", ts(2000, 0)), vec![0]), + // One nanosecond past a millisecond is not rounded away. + (b.greater_or_equal("ts", ts(2000, 1)), vec![2, 3]), + (b.between("ts", ts(1500, 0), ts(3000, 0)), vec![1, 2]), + (b.not_between("ts", ts(1500, 0), ts(3000, 0)), vec![0, 3]), + (b.is_not_in("ts", vec![ts(1000, 0)]), vec![1, 2, 3]), + ]; + for (predicate, expected) in cases { + let predicate = predicate.unwrap(); + assert_eq!( + kept(std::slice::from_ref(&predicate)), + expected, + "{predicate}" + ); + } + } + + #[test] + fn test_every_nan_sorts_above_every_number() { + let fields = vec![ + field(0, "k", DataType::Int(IntType::with_nullable(false))), + field(1, "d", DataType::Double(DoubleType::new())), + field(2, "f", DataType::Float(FloatType::new())), + ]; + // Arrow's total order alone would put the sign-bit NaN below every number. + let rows = batch( + &fields, + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + Arc::new(Float64Array::from(vec![-f64::NAN, f64::NAN, -1.0, 1.0])), + Arc::new(Float32Array::from(vec![-f32::NAN, f32::NAN, -1.0, 1.0])), + ], + ); + for column in ["d", "f"] { + for (function, literals, expected) in [ + ("LESS_THAN", "[0.0]", vec![2]), + ("GREATER_THAN", "[0.0]", vec![0, 1, 3]), + ("GREATER_THAN", "[1.0]", vec![0, 1]), + ("NOT_BETWEEN", "[-2.0,2.0]", vec![0, 1]), + ] { + let json = leaf(function, column, literals); + let rules = parse(&[&json], &[], &fields).unwrap(); + let out = filter_batch(&rows, &rules.filters, &fields, &fields).unwrap(); + assert_eq!(ids(&out), expected, "{json}"); + } + } + } + + #[test] + fn test_filters_combine_with_and_or() { + let fields = typed_fields(); + let i_is = |value: &str| leaf("EQUAL", "i", &format!("[{value}]")); + let cases = [ + ( + vec![compound( + "OR", + &[i_is("1"), leaf("EQUAL", "s", r#"["banana"]"#)], + )], + vec![0, 2], + ), + ( + vec![compound( + "AND", + &[ + leaf("GREATER_OR_EQUAL", "b", "[20]"), + leaf("EQUAL", "flag", "[true]"), + ], + )], + vec![2], + ), + // A null fails NOT_EQUAL but passes IS_NULL. + ( + vec![compound( + "OR", + &[leaf("NOT_EQUAL", "i", "[2]"), leaf("IS_NULL", "s", "[]")], + )], + vec![0, 2, 3, 4], + ), + ( + vec![compound( + "AND", + &[ + compound("OR", &[i_is("1"), i_is("4")]), + leaf("NOT_EQUAL", "s", r#"["apple"]"#), + ], + )], + vec![3], + ), + // The server's entries are ANDed. + ( + vec![ + leaf("GREATER_THAN", "i", "[1]"), + leaf("EQUAL", "flag", "[true]"), + ], + vec![2], + ), + (vec![constant("TRUE")], vec![0, 1, 2, 3, 4]), + (vec![constant("FALSE"), leaf("IS_NULL", "i", "[]")], vec![]), + ]; + for (filters, expected) in cases { + let filters: Vec<&str> = filters.iter().map(String::as_str).collect(); + let rules = parse(&filters, &[], &fields).unwrap(); + assert_eq!(kept(&rules.filters), expected, "{filters:?}"); + } + assert_eq!(kept(&[]), vec![0, 1, 2, 3, 4]); + } + + #[test] + fn test_filter_refuses_what_it_cannot_evaluate() { + let fields = typed_fields(); + let rows = typed_batch(); + let json = leaf("EQUAL", "i", "[1]"); + let filters = parse(&[&json], &[], &fields).unwrap().filters; + + // Carried under its id and its name, or not at all. + for carried in [ + field(1, "renamed", DataType::Int(IntType::new())), + field(99, "i", DataType::Int(IntType::new())), + ] { + let mut batch_fields = fields.clone(); + batch_fields[1] = carried; + let message = refusal(filter_batch(&rows, &filters, &fields, &batch_fields)); + assert!(message.contains("does not carry 'i'"), "{message}"); + } + let narrow = rows.project(&[0, 2]).unwrap(); + let narrow_fields = [fields[0].clone(), fields[2].clone()]; + let message = refusal(filter_batch(&narrow, &filters, &fields, &narrow_fields)); + assert!(message.contains("does not carry 'i'"), "{message}"); + // Any position will do. + let reordered = rows.project(&[1, 0]).unwrap(); + let reordered_fields = [fields[1].clone(), fields[0].clone()]; + let out = filter_batch(&reordered, &filters, &fields, &reordered_fields).unwrap(); + assert_eq!(out.num_rows(), 1); + + let mistyped = Predicate::Leaf { + column: "i".to_string(), + index: 1, + data_type: DataType::Int(IntType::new()), + op: PredicateOperator::Eq, + literals: vec![Datum::String("1".to_string())], + }; + for (filter, reason) in [ + (Predicate::Not(Box::new(filters[0].clone())), "has a NOT"), + (Predicate::And(Vec::new()), "empty AND/OR"), + (mistyped, "cannot evaluate"), + ] { + let message = refusal(filter_batch(&rows, &[filter], &fields, &fields)); + assert!(message.contains(reason), "{message}"); + } + let message = refusal(filter_batch(&rows, &filters, &fields[..1], &fields)); + assert!(message.contains("unknown column"), "{message}"); + } + + #[test] + fn test_a_zero_column_batch_keeps_its_row_count() { + let fields = rule_fields(); + let options = RecordBatchOptions::new().with_row_count(Some(3)); + let empty = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), Vec::new(), &options) + .unwrap(); + for (filter, rows) in [(constant("TRUE"), 3), (constant("FALSE"), 0)] { + let rules = parse(&[&filter], &[], &fields).unwrap(); + let out = filter_batch(&empty, &rules.filters, &fields, &[]).unwrap(); + assert_eq!(out.num_rows(), rows, "{filter}"); + } + assert_eq!( + filter_batch(&empty, &[], &fields, &[]).unwrap().num_rows(), + 3 + ); + } +} diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index ab501eebe..188cc0101 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -540,6 +540,15 @@ impl DataSplit { self.query_auth_grant.as_ref() } + /// Whether the server attached rules to this split's read: its row counts + /// and column statistics then describe rows the read will not return as + /// they are. + pub fn has_query_auth_rules(&self) -> bool { + self.query_auth_grant + .as_ref() + .is_some_and(|grant| !grant.is_unrestricted()) + } + pub fn snapshot_id(&self) -> i64 { self.snapshot_id } @@ -658,8 +667,18 @@ impl DataSplit { /// [`DataFileMeta::ROW_COUNT_UNKNOWN`] — no arithmetic over a placeholder /// produces a number a caller may trust. /// + /// Also `None` under a query-auth row filter, which drops rows the count + /// includes, as Java `QueryAuthSplit` does. + /// /// Reference: [DataSplit.mergedRowCount()](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java#L133) pub fn merged_row_count(&self) -> Option { + if self + .query_auth_grant + .as_ref() + .is_some_and(|grant| !grant.rules().filters.is_empty()) + { + return None; + } if let Some(ranges) = &self.row_ranges { return Some(ranges.iter().map(RowRange::count).sum()); } @@ -1606,6 +1625,24 @@ mod tests { assert_eq!(s.merged_row_count(), Some(15)); } + #[tokio::test] + async fn test_merged_row_count_is_unknown_under_a_row_filter() { + let table = crate::table::rest_query_auth_table().await; + let grant = |filters| { + Some(Arc::new(crate::table::query_auth::QueryAuthGrant::new( + table.query_auth_session().unwrap(), + None, + crate::table::query_auth::Rules { filters }, + ))) + }; + let s = || split(vec![file("a", 10, None)], true); + assert_eq!(s().planned(grant(Vec::new())).merged_row_count(), Some(10)); + let filtered = s().planned(grant(vec![crate::spec::Predicate::AlwaysTrue])); + assert_eq!(filtered.merged_row_count(), None); + // The physical count stays, as Java's `rowCount` does. + assert_eq!(filtered.row_count(), 10); + } + /// A placeholder row count must surface as "unknown", not as a number. /// Reporting `Some(0)` here lets an engine answer `COUNT(*)` with 0 for a /// split that holds data. diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index ee0d3a0fc..1c6fd65bf 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -312,22 +312,6 @@ impl<'a> TableRead<'a> { } } -/// Every leaf's column name, system columns included. -fn collect_leaf_column_names(predicate: &Predicate, out: &mut std::collections::HashSet) { - match predicate { - Predicate::Leaf { column, .. } => { - out.insert(column.clone()); - } - Predicate::And(children) | Predicate::Or(children) => { - children - .iter() - .for_each(|child| collect_leaf_column_names(child, out)); - } - Predicate::Not(inner) => collect_leaf_column_names(inner, out), - Predicate::AlwaysTrue | Predicate::AlwaysFalse => {} - } -} - #[derive(Debug, Clone)] struct PaimonTableRead<'a> { table: &'a Table, @@ -508,7 +492,14 @@ impl<'a> PaimonTableRead<'a> { data_splits: &[DataSplit], ) -> crate::Result { // Streaming primary-key splits are read raw below, so decide here for all. - self.ensure_authorized_by_splits(&self.table.schema.core_options(), data_splits)?; + if self + .ensure_authorized_by_splits(&self.table.schema.core_options(), data_splits)? + .is_some() + { + return Err(super::query_auth::unsupported( + "a row-kind read cannot apply a row filter or column masking", + )); + } let schema = audit_schema_for_read_type(&self.read_type, false)?; let (streaming, materialized): (Vec<_>, Vec<_>) = data_splits .iter() @@ -889,13 +880,13 @@ impl<'a> PaimonTableRead<'a> { reader.read(splits) } - /// Reads only splits carrying a grant that the server imposed nothing; - /// nothing is fetched here, so a split without one fails closed. + /// Every split must carry this handle's grant; returns it when it restricts + /// the read. Nothing is fetched here, so a split without one fails closed. fn ensure_authorized_by_splits( &self, core_options: &CoreOptions, data_splits: &[DataSplit], - ) -> crate::Result<()> { + ) -> crate::Result>> { // Unconditional: unrelated to query-auth. core_options.ensure_type_paimon_served(&self.table.identifier().full_name())?; // Decided at plan time, as in Java: a split predating the option, or built @@ -903,7 +894,7 @@ impl<'a> PaimonTableRead<'a> { let required = core_options.query_auth_enabled() || data_splits.iter().any(|s| s.query_auth_required()); if !required { - return Ok(()); + return Ok(None); } // Only the catalog mints a session, so a handle without one holds no grant. if self.table.query_auth_session().is_none() { @@ -912,10 +903,7 @@ impl<'a> PaimonTableRead<'a> { )); } // The read's own scope: a caller can plan clean, then read differently. - let mut filter_columns = std::collections::HashSet::new(); - for predicate in &self.data_predicates { - collect_leaf_column_names(predicate, &mut filter_columns); - } + let filter_columns = super::query_auth::leaf_names(&self.data_predicates); super::query_auth::reject_system_columns( self.read_type .iter() @@ -928,7 +916,9 @@ impl<'a> PaimonTableRead<'a> { &self.read_type, self.table.schema().fields(), )?; - // Per split, as Java binds one `QueryAuthSplit` each. + // Per split, as Java binds one `QueryAuthSplit` each; one read applies one + // rule set, so a restricted split's grant must cover them all. + let mut restricted: Option<&Arc> = None; for split in data_splits { let Some(grant) = split.query_auth_grant() else { return Err(super::query_auth::unsupported( @@ -942,21 +932,124 @@ impl<'a> PaimonTableRead<'a> { scan", )); } - if !grant.is_unrestricted() { + // The server ruled on these columns only, whether or not it set rules. + if let Some(select) = grant.select() { + if let Some(outside) = self + .read_type + .iter() + .map(|f| f.name()) + .chain(filter_columns.iter().map(String::as_str)) + .find(|name| !select.iter().any(|s| s == name)) + { + return Err(super::query_auth::unsupported(&format!( + "'{outside}' is outside the columns this plan was authorized for; plan \ + with it" + ))); + } + } + if !grant.is_unrestricted() && restricted.is_none() { + restricted = Some(grant); + } + } + if let Some(grant) = restricted { + if data_splits + .iter() + .any(|s| s.query_auth_grant().is_none_or(|g| **g != **grant)) + { return Err(super::query_auth::unsupported( - "this client cannot apply a row filter or column masking, so it refuses \ - rather than return unfiltered rows", + "the splits were planned under different rules; re-plan the scan", )); } } - Ok(()) + Ok(restricted.cloned()) } /// Returns an [`ArrowRecordBatchStream`]. pub fn to_arrow(&self, data_splits: &[DataSplit]) -> crate::Result { - let has_primary_keys = !self.table.schema.primary_keys().is_empty(); let core_options = self.table.schema.core_options(); - self.ensure_authorized_by_splits(&core_options, data_splits)?; + match self.ensure_authorized_by_splits(&core_options, data_splits)? { + Some(grant) => self.read_restricted(data_splits, &grant), + None => self.read_splits(data_splits, &core_options), + } + } + + /// Reads what the row filter needs, filters on stored values and projects + /// back (Java `doAuth`). + fn read_restricted( + &self, + data_splits: &[DataSplit], + grant: &super::query_auth::QueryAuthGrant, + ) -> crate::Result { + use super::query_auth::{filter_batch, unsupported}; + + let schema_fields = self.table.schema().fields().to_vec(); + let rules = grant.rules(); + let index_of = |field: &DataField| schema_fields.iter().position(|s| s.id() == field.id()); + let needed = rules.filter_columns(); + // A partly projected column would feed the filter a partial value (Java + // `validateReadType`). + for field in &self.read_type { + if let Some(index) = index_of(field) { + if needed.contains(&index) && field.data_type() != schema_fields[index].data_type() + { + return Err(unsupported(&format!( + "the server's row filter reads '{}', which the read projects only in part", + field.name() + ))); + } + } + } + + let mut physical = self.read_type.clone(); + let mut needed: Vec = needed.into_iter().collect(); + needed.sort_unstable(); + for index in needed { + let field = &schema_fields[index]; + if !physical.iter().any(|f| f.id() == field.id()) { + physical.push(field.clone()); + } + } + let mut inner = self.clone(); + inner.read_type = physical.clone(); + inner.limit = None; + let stream = inner.read_splits(data_splits, &self.table.schema.core_options())?; + + let filters = rules.filters.clone(); + let projection: Vec = (0..self.read_type.len()).collect(); + let stream = stream.map(move |batch| { + let batch = batch?; + let names_match = batch.num_columns() == physical.len() + && batch + .schema() + .fields() + .iter() + .zip(&physical) + .all(|(column, field)| column.name() == field.name()); + if !names_match { + return Err(unsupported( + "the read returned columns the rules cannot address", + )); + } + let batch = filter_batch(&batch, &filters, &schema_fields, &physical)?; + batch + .project(&projection) + .map_err(|e| crate::Error::DataInvalid { + message: format!("failed to project the authorized batch: {e}"), + source: Some(Box::new(e)), + }) + }); + Ok(match self.limit { + None => Box::pin(stream), + Some(limit) => limit_rows(Box::pin(stream), limit), + }) + } + + fn read_splits( + &self, + data_splits: &[DataSplit], + core_options: &CoreOptions<'_>, + ) -> crate::Result { + let has_primary_keys = !self.table.schema.primary_keys().is_empty(); let merge_engine = core_options.merge_engine()?; // Route supported PK merge engines through the split-aware reader. @@ -971,11 +1064,11 @@ impl<'a> PaimonTableRead<'a> { | MergeEngine::Aggregation ) { - return self.read_pk_with_blob(data_splits, &core_options); + return self.read_pk_with_blob(data_splits, core_options); } if core_options.data_evolution_enabled() { - self.read_with_evolution(data_splits, &core_options) + self.read_with_evolution(data_splits, core_options) } else { self.read_raw(data_splits) } @@ -1176,6 +1269,26 @@ impl<'a> PaimonTableRead<'a> { } } +/// Stops once `limit` rows are out, without polling for more. +fn limit_rows(stream: ArrowRecordBatchStream, limit: usize) -> ArrowRecordBatchStream { + Box::pin(stream::unfold( + (stream, 0usize), + move |(mut inner, emitted)| async move { + if emitted >= limit { + return None; + } + match inner.next().await? { + Err(e) => Some((Err(e), (inner, limit))), + Ok(batch) => { + let batch = batch.slice(0, batch.num_rows().min(limit - emitted)); + let emitted = emitted + batch.num_rows(); + Some((Ok(batch), (inner, emitted))) + } + } + }, + )) +} + fn prepend_insert_row_kind_stream( stream: ArrowRecordBatchStream, schema: Arc, @@ -2351,13 +2464,17 @@ mod tests { fn grant_for(table: &Table, restricted: bool) -> crate::table::query_auth::QueryAuthGrant { crate::table::query_auth::QueryAuthGrant::new( - crate::api::AuthTableQueryResponse { - filter: restricted.then(|| vec!["{}".to_string()]), - column_masking: None, - }, table .query_auth_session() .expect("a catalog-loaded table has a session"), + None, + crate::table::query_auth::Rules { + filters: if restricted { + vec![crate::spec::Predicate::AlwaysTrue] + } else { + Vec::new() + }, + }, ) } @@ -2520,17 +2637,51 @@ mod tests { } #[tokio::test] - async fn test_restricted_grant_on_a_split_refuses_the_read() { + async fn test_a_restricted_grant_must_cover_every_split_of_the_read() { let table = crate::table::rest_query_auth_table().await; let read = TableRead::new(&table, table.schema.fields().to_vec(), Vec::new()); - let split = split_with_grant(Some(grant_for(&table, true))); - assert!( - matches!( - read.to_arrow(&[split]), - Err(crate::Error::Unsupported { ref message }) if message.contains("query-auth.enabled") - ), - "a row filter this client cannot apply must refuse the read" - ); + let restricted = split_with_grant(Some(grant_for(&table, true))); + assert!(read.to_arrow(std::slice::from_ref(&restricted)).is_ok()); + let unrestricted = split_with_grant(Some(grant_for(&table, false))); + for splits in [ + [restricted.clone(), unrestricted.clone()], + [unrestricted, restricted.clone()], + ] { + assert!( + matches!(read.to_arrow(&splits), Err(crate::Error::Unsupported { ref message }) + if message.contains("different rules")), + "one rule set per read, whatever the order" + ); + } + // Streaming primary-key splits are read raw. + assert!(matches!( + read.to_arrow_with_row_kind(&[restricted]), + Err(crate::Error::Unsupported { ref message }) if message.contains("row-kind") + )); + } + + #[tokio::test] + async fn test_a_read_stays_inside_the_authorized_columns() { + let table = crate::table::rest_query_auth_table().await; + let fields = table.schema.fields().to_vec(); + // Without rules too: a column permission scoped the grant all the same. + for filters in [vec![Predicate::AlwaysTrue], Vec::new()] { + // Planned with no columns, as `COUNT(*)` is. + let grant = crate::table::query_auth::QueryAuthGrant::new( + table.query_auth_session().unwrap(), + Some(Vec::new()), + crate::table::query_auth::Rules { filters }, + ); + let split = split_with_grant(Some(grant)); + let narrow = TableRead::new(&table, Vec::new(), Vec::new()); + assert!(narrow.to_arrow(std::slice::from_ref(&split)).is_ok()); + let wide = TableRead::new(&table, fields.clone(), Vec::new()); + assert!( + matches!(wide.to_arrow(&[split]), Err(crate::Error::Unsupported { ref message }) + if message.contains("outside the columns")), + "the server ruled on the planned columns only" + ); + } } #[tokio::test] diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 6205cd116..b17b5e080 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1501,32 +1501,42 @@ impl<'a> PaimonTableScan<'a> { /// `scan.snapshot-id` / `scan.tag-name` handling. pub async fn plan(&self) -> crate::Result { let grant = self.authorize_query().await?; - self.validate_shard_strategy()?; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; - let plan = match super::time_travel::resolve_snapshot(self.table).await? { + let restricted = grant + .as_deref() + .filter(|grant| !grant.is_unrestricted()) + .map(|grant| self.restricted_by(grant)); + let scan = restricted.as_ref().unwrap_or(self); + scan.validate_shard_strategy()?; + let data_evolution_read_field_ids = scan.projected_read_field_ids()?; + let plan = match super::time_travel::resolve_snapshot(scan.table).await? { Some(snapshot) => { - self.plan_snapshot(snapshot, data_evolution_read_field_ids.as_ref(), None) + scan.plan_snapshot(snapshot, data_evolution_read_field_ids.as_ref(), None) .await? } None => Plan::new(Vec::new()), }; - self.check_planned_files(&plan, grant.is_some()).await?; + scan.check_planned_files(&plan, grant.is_some()).await?; Ok(plan.planned(grant)) } /// Plan the full scan and return metadata-pruning trace counters. pub async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { let grant = self.authorize_query().await?; - self.validate_shard_strategy()?; + let restricted = grant + .as_deref() + .filter(|grant| !grant.is_unrestricted()) + .map(|grant| self.restricted_by(grant)); + let scan = restricted.as_ref().unwrap_or(self); + scan.validate_shard_strategy()?; let mut trace = ScanTrace { - limit: self.limit, + limit: scan.limit, ..Default::default() }; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; - let plan = match super::time_travel::resolve_snapshot(self.table).await? { + let data_evolution_read_field_ids = scan.projected_read_field_ids()?; + let plan = match super::time_travel::resolve_snapshot(scan.table).await? { Some(snapshot) => { trace.snapshot_id = Some(snapshot.id()); - let plan = self + let plan = scan .plan_snapshot( snapshot, data_evolution_read_field_ids.as_ref(), @@ -1538,10 +1548,67 @@ impl<'a> PaimonTableScan<'a> { } None => Plan::new(Vec::new()), }; - self.check_planned_files(&plan, grant.is_some()).await?; + scan.check_planned_files(&plan, grant.is_some()).await?; Ok((plan.planned(grant), trace)) } + /// The scan a restricted grant plans: a limit cannot count rows the row + /// filter drops. + fn restricted_by(&self, grant: &super::query_auth::QueryAuthGrant) -> Self { + let rules = grant.rules(); + let mut scan = self.clone(); + if !rules.filters.is_empty() { + scan.limit = scan.limit.filter(|limit| *limit == 0); + } + // Column-slice pruning must keep the files holding the filter's columns. + if let Some(ids) = scan.projected_read_field_ids.as_mut() { + let fields = self.table.schema().fields(); + ids.extend( + rules + .filter_columns() + .into_iter() + .filter_map(|i| fields.get(i)) + .map(|f| f.id()), + ); + } + scan + } + + /// Java's `select`: the projection plus what the filters read; `None` (the + /// whole table) without a projection. + fn query_auth_select(&self) -> Option> { + let ids = self.projected_read_field_ids.as_ref()?; + let mut select: Vec = self + .table + .schema() + .fields() + .iter() + .filter(|f| ids.contains(&f.id())) + .map(|f| f.name().to_string()) + .collect(); + for key in self.partition_filter_columns() { + if !select.contains(&key) { + select.push(key); + } + } + Some(select) + } + + /// Columns the caller's partition filter reads; a pinned partition set reads them all. + fn partition_filter_columns(&self) -> Vec { + match &self.partition_filter { + None => Vec::new(), + Some(PartitionFilter::Predicate(predicate)) => { + super::query_auth::leaf_names(std::slice::from_ref(predicate)) + .into_iter() + .collect() + } + Some(PartitionFilter::PartitionSet { .. }) => { + self.table.schema().partition_keys().to_vec() + } + } + } + /// The grant predates the manifest read, so the table can have been /// re-created in between. Also refuses stats the schema no longer covers. async fn check_planned_files(&self, plan: &Plan, query_auth: bool) -> crate::Result<()> { @@ -1592,14 +1659,10 @@ impl<'a> PaimonTableScan<'a> { super::query_auth::reject_system_columns([ROW_ID_FIELD_NAME])?; } - let grant = self.table.authorize_read(query_auth).await?; - // A plan already answers COUNT/MIN/MAX from row counts and bounds. - if grant.as_ref().is_some_and(|g| !g.is_unrestricted()) { - return Err(super::query_auth::unsupported( - "a plan already carries file paths, row counts and column bounds that a row \ - filter or column masking must not expose", - )); - } + let grant = self + .table + .authorize_read(query_auth, self.query_auth_select()) + .await?; Ok(grant) } diff --git a/crates/paimon/tests/mock_server.rs b/crates/paimon/tests/mock_server.rs index 68dadf346..fd2082ac6 100644 --- a/crates/paimon/tests/mock_server.rs +++ b/crates/paimon/tests/mock_server.rs @@ -89,6 +89,8 @@ struct MockState { drop_policy_error: Option, auth_responses: HashMap, column_auth: HashMap>, + /// The `select` of every auth request, per `"{db}.{table}"`. + auth_selects: HashMap>>>, uuid_after_auth: HashMap, uuid_after_calls: HashMap, /// ECS metadata role name (for token loader testing) @@ -871,8 +873,12 @@ impl RESTServer { Extension(state): Extension>, Json(request): Json, ) -> impl IntoResponse { - let s = state.inner.lock().unwrap(); + let mut s = state.inner.lock().unwrap(); let key = format!("{db}.{table}"); + s.auth_selects + .entry(key.clone()) + .or_default() + .push(request.select.clone()); // Mirrors the reference server: a null select means the real schema // fields, and any column outside the grant denies the query. @@ -938,6 +944,14 @@ impl RESTServer { s.column_auth.insert(format!("{database}.{table}"), columns); } + pub fn auth_selects(&self, database: &str, table: &str) -> Vec>> { + let s = self.inner.lock().unwrap(); + s.auth_selects + .get(&format!("{database}.{table}")) + .cloned() + .unwrap_or_default() + } + /// Handle DELETE /databases/:db/tables/:table - drop a table. pub async fn drop_table( Path((db, table)): Path<(String, String)>, diff --git a/crates/paimon/tests/rest_catalog_test.rs b/crates/paimon/tests/rest_catalog_test.rs index 24f217c3c..0ee96486d 100644 --- a/crates/paimon/tests/rest_catalog_test.rs +++ b/crates/paimon/tests/rest_catalog_test.rs @@ -2802,27 +2802,64 @@ fn assert_drifted(err: paimon::Error, what: &str) { ); } -fn restricted() -> paimon::api::AuthTableQueryResponse { +/// A row filter of `filters` (ANDed) and `masks` by column, as the server sends them. +fn rules( + filters: &[serde_json::Value], + masks: &[(&str, serde_json::Value)], +) -> paimon::api::AuthTableQueryResponse { paimon::api::AuthTableQueryResponse { - filter: Some(vec!["{\"field\":\"id\"}".to_string()]), - column_masking: None, + filter: (!filters.is_empty()).then(|| filters.iter().map(|f| f.to_string()).collect()), + column_masking: (!masks.is_empty()).then(|| { + masks + .iter() + .map(|(column, mask)| (column.to_string(), mask.to_string())) + .collect() + }), } } -#[tokio::test] -async fn test_query_auth_restricted_user_is_refused_at_plan_time() { - let g = guarded("restricted", &["id"]).await; - g.ctx - .server - .set_auth_response("default", "restricted", restricted()); - - assert_refused( - plan_err( - &g.table, - "a restricted user must be refused before a plan exists", - ) - .await, - ); +/// Java `LeafPredicate` JSON over an INT column. +fn int_leaf(index: usize, column: &str, function: &str, literal: i32) -> serde_json::Value { + serde_json::json!({ + "kind": "LEAF", + "transform": { + "name": "FIELD_REF", + "fieldRef": {"index": index, "name": column, "type": "INT"}, + }, + "function": function, + "literals": [literal], + }) +} + +#[tokio::test] +async fn test_query_auth_rules_this_client_cannot_parse_fail_closed_at_planning() { + let g = guarded("unparseable", &["id"]).await; + let widened = serde_json::json!({ + "name": "CAST", + "fieldRef": {"index": 0, "name": "id", "type": "INT"}, + "type": "BIGINT", + }); + for response in [ + // Not a predicate, or nothing at all. + rules(&[serde_json::json!({"field": "id"})], &[]), + paimon::api::AuthTableQueryResponse { + filter: Some(vec![String::new()]), + column_masking: None, + }, + // Functions and transforms this client does not know. + rules(&[int_leaf(0, "id", "REGEXP", 1)], &[]), + rules(&[], &[("id", serde_json::json!({"name": "invalid"}))]), + // Columns the table does not have. + rules(&[int_leaf(0, "missing", "EQUAL", 1)], &[]), + rules(&[], &[("missing", serde_json::json!({"name": "NULL"}))]), + // A mask must keep its column's type. + rules(&[], &[("id", widened)]), + ] { + g.ctx + .server + .set_auth_response("default", "unparseable", response.clone()); + assert_refused(plan_err(&g.table, &format!("{response:?}")).await); + } } #[tokio::test] @@ -3075,6 +3112,363 @@ async fn test_query_auth_refuses_an_assembled_handle() { } } +fn people(options: &[(&str, &str)]) -> Schema { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())); + for (key, value) in options { + builder = builder.option(*key, *value); + } + builder.build().unwrap() +} + +fn people_batch(rows: &[(i32, &str)]) -> RecordBatch { + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.0))), + Arc::new(StringArray::from_iter_values(rows.iter().map(|r| r.1))), + ], + ) + .unwrap() +} + +const NAMES: [&str; 10] = [ + "alice", "bob", "carol", "dave", "erin", "frank", "grace", "heidi", "ivan", "judy", +]; + +/// Ids `ids`, each with its name from [`NAMES`]. +fn people_rows(ids: std::ops::RangeInclusive) -> RecordBatch { + let rows: Vec<(i32, &str)> = ids.map(|id| (id, NAMES[id as usize - 1])).collect(); + people_batch(&rows) +} + +/// Writes one commit per batch through a filesystem catalog, then serves the +/// files as a `query-auth.enabled` table. +async fn written( + name: &str, + schema: impl Fn(&[(&str, &str)]) -> Schema, + commits: Vec, +) -> Guarded { + let tmp = tempfile::tempdir().unwrap(); + let mut fs_options = Options::new(); + fs_options.set( + CatalogOptions::WAREHOUSE, + format!("file://{}", tmp.path().display()), + ); + let fs_catalog = FileSystemCatalog::new(fs_options).unwrap(); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let identifier = Identifier::new("default", name); + fs_catalog + .create_table(&identifier, schema(&[]), false) + .await + .unwrap(); + let plain = fs_catalog.get_table(&identifier).await.unwrap(); + for batch in commits { + write_batch(&plain, batch, "query-auth-fixture").await; + } + + let ctx = setup_catalog(vec!["default"]).await; + ctx.server + .add_table_with_schema("default", name, schema(GUARDED), plain.location()); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + Guarded { + ctx, + table, + identifier, + _tmp: tmp, + } +} + +async fn read_all(builder: &paimon::table::ReadBuilder<'_>) -> paimon::Result> { + let plan = builder.new_scan().plan().await?; + builder + .new_read()? + .to_arrow(plan.splits())? + .try_collect() + .await +} + +fn column<'b, T: 'static>(batch: &'b RecordBatch, name: &str) -> &'b T { + batch + .column_by_name(name) + .unwrap_or_else(|| panic!("no column '{name}'")) + .as_any() + .downcast_ref::() + .unwrap() +} + +fn ids(batches: &[RecordBatch]) -> Vec { + let mut ids: Vec = batches + .iter() + .flat_map(|b| column::(b, "id").values().to_vec()) + .collect(); + ids.sort_unstable(); + ids +} + +/// `(id, name)` sorted by id; a null name is `None`. +fn id_names(batches: &[RecordBatch]) -> Vec<(i32, Option)> { + let mut rows: Vec<_> = batches + .iter() + .flat_map(|b| { + let ids = column::(b, "id"); + let names = column::(b, "name"); + (0..b.num_rows()) + .map(|r| { + ( + ids.value(r), + names.is_valid(r).then(|| names.value(r).to_string()), + ) + }) + .collect::>() + }) + .collect(); + rows.sort(); + rows +} + +#[track_caller] +fn assert_forbidden(err: paimon::Error) { + assert!( + matches!( + err, + paimon::Error::RestApi { + source: paimon::api::RestError::Forbidden { .. } + } + ), + "{err:?}" + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_row_filter_returns_only_the_rows_it_admits() { + // One split per commit, so a pushed-down limit could stop after the first. + let g = written( + "filtered", + |options| people(&[options, &[("source.split.target-size", "1b")]].concat()), + vec![people_rows(1..=5), people_rows(6..=10)], + ) + .await; + g.ctx.server.set_auth_response( + "default", + "filtered", + rules(&[int_leaf(0, "id", "GREATER_THAN", 6)], &[]), + ); + + let rows = id_names(&read_all(&g.table.new_read_builder()).await.unwrap()); + let expected: Vec<_> = (7..=10) + .map(|id| (id, Some(NAMES[id as usize - 1].to_string()))) + .collect(); + assert_eq!(rows, expected); + + // What COUNT(*) reads: no columns, only the admitted rows. + let mut counted = g.table.new_read_builder(); + counted.with_projection(&[]).unwrap(); + let batches = read_all(&counted).await.unwrap(); + assert!(batches.iter().all(|b| b.num_columns() == 0)); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 4); + + let mut limited = g.table.new_read_builder(); + limited.with_limit(2); + let plan = limited.new_scan().plan().await.unwrap(); + assert!(plan.splits().iter().all(|s| s.has_query_auth_rules())); + assert_eq!( + plan.splits().len(), + 2, + "the row filter must disable limit pushdown" + ); + let batches: Vec = limited + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let limited_ids = ids(&batches); + assert_eq!(limited_ids.len(), 2, "{limited_ids:?}"); + assert!(limited_ids.iter().all(|id| *id > 6), "{limited_ids:?}"); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_asks_about_the_projection_and_the_columns_filters_read() { + let g = written("scoped", people, vec![people_rows(1..=3)]).await; + g.ctx + .server + .set_column_auth("default", "scoped", vec!["id".to_string()]); + + let mut id_only = g.table.new_read_builder(); + id_only.with_projection(&["id"]).unwrap(); + assert_eq!(ids(&read_all(&id_only).await.unwrap()), vec![1, 2, 3]); + + let mut with_name = g.table.new_read_builder(); + with_name.with_projection(&["id", "name"]).unwrap(); + assert_forbidden(read_all(&with_name).await.unwrap_err()); + + let mut filtered = g.table.new_read_builder(); + filtered.with_projection(&["id"]).unwrap().with_filter( + PredicateBuilder::new(g.table.schema().fields()) + .equal("name", Datum::String("bob".to_string())) + .unwrap(), + ); + assert_forbidden(read_all(&filtered).await.unwrap_err()); + + assert_forbidden(read_all(&g.table.new_read_builder()).await.unwrap_err()); + + // The order of a select is not part of what it asks. + let asked: Vec>> = g + .ctx + .server + .auth_selects("default", "scoped") + .into_iter() + .map(|select| select.map(|columns| columns.into_iter().collect())) + .collect(); + let select = |columns: &[&str]| Some(columns.iter().map(|c| c.to_string()).collect()); + assert_eq!( + asked, + vec![ + select(&["id"]), + select(&["id", "name"]), + select(&["id", "name"]), + None + ] + ); +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_refuses_a_read_wider_than_its_plan() { + let g = written("narrow", people, vec![people_rows(1..=10)]).await; + let mut planned = g.table.new_read_builder(); + planned.with_projection(&["id"]).unwrap(); + let mut filtered = g.table.new_read_builder(); + filtered.with_projection(&["id"]).unwrap().with_filter( + PredicateBuilder::new(g.table.schema().fields()) + .equal("name", Datum::String("bob".to_string())) + .unwrap(), + ); + let wider = [g.table.new_read_builder(), filtered]; + + // Scoped once under a row filter, once by a column permission alone. + g.ctx.server.set_auth_response( + "default", + "narrow", + rules(&[int_leaf(0, "id", "GREATER_THAN", 6)], &[]), + ); + let under_rules = planned.new_scan().plan().await.unwrap(); + g.ctx + .server + .set_auth_response("default", "narrow", rules(&[], &[])); + g.ctx + .server + .set_column_auth("default", "narrow", vec!["id".to_string()]); + let under_permission = planned.new_scan().plan().await.unwrap(); + + for (plan, admitted) in [(under_rules, 4), (under_permission, 10)] { + let read: Vec = planned + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(ids(&read).len(), admitted); + for builder in &wider { + let Err(err) = builder.new_read().unwrap().to_arrow(plan.splits()) else { + panic!("'name' was never authorized"); + }; + assert!( + matches!(err, paimon::Error::Unsupported { ref message } + if message.contains("query-auth.enabled") + && message.contains("outside the columns")), + "{err:?}" + ); + } + } +} + +/// A primary-key table whose second commit updates keys 1 and 2. +async fn scores() -> Guarded { + let schema = |options: &[(&str, &str)]| { + let mut builder = Schema::builder() + .column("id", DataType::Int(IntType::with_nullable(false))) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("score", DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1"); + for (key, value) in options { + builder = builder.option(*key, *value); + } + builder.build().unwrap() + }; + let batch = |rows: &[(i32, &str, i32)]| { + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("score", ArrowDataType::Int32, true), + ])), + vec![ + Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.0))), + Arc::new(StringArray::from_iter_values(rows.iter().map(|r| r.1))), + Arc::new(Int32Array::from_iter_values(rows.iter().map(|r| r.2))), + ], + ) + .unwrap() + }; + written( + "scores", + schema, + vec![ + batch(&[(1, "alice", 10), (2, "bob", 50), (3, "carol", 90)]), + batch(&[(1, "alicia", 80), (2, "bobby", 20)]), + ], + ) + .await +} + +fn score_rows(batches: &[RecordBatch]) -> Vec<(i32, String, i32)> { + let mut rows: Vec<(i32, String, i32)> = batches + .iter() + .flat_map(|b| { + let ids = column::(b, "id"); + let names = column::(b, "name"); + let scores = column::(b, "score"); + (0..b.num_rows()) + .map(|r| (ids.value(r), names.value(r).to_string(), scores.value(r))) + .collect::>() + }) + .collect(); + rows.sort(); + rows +} + +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_filters_the_merged_primary_key_row() { + let g = scores().await; + g.ctx.server.set_auth_response( + "default", + "scores", + rules(&[int_leaf(2, "score", "GREATER_THAN", 30)], &[]), + ); + // Filtering each commit before the merge would bring back bob's first row. + assert_eq!( + score_rows(&read_all(&g.table.new_read_builder()).await.unwrap()), + vec![(1, "alicia".to_string(), 80), (3, "carol".to_string(), 90)] + ); +} + #[tokio::test] async fn test_rest_catalog_manages_permissions_end_to_end() { let ctx = setup_catalog(vec!["default"]).await; diff --git a/docs/src/sql.md b/docs/src/sql.md index 259b2b2ac..630587044 100644 --- a/docs/src/sql.md +++ b/docs/src/sql.md @@ -1537,6 +1537,19 @@ When the following conditions are met, `COUNT(*)` retrieves exact row counts dir - No LIMIT clause - Filter predicates only involve partition columns (Exact level) +### Query Authorization + +When a table in a REST catalog has `'query-auth.enabled' = 'true'`, planning a +query asks the catalog what the current user may read, and the row filters it +returns (see `create_policy`) are applied to the result: + +- Rows are filtered on their stored values. +- A user the catalog also masks columns for is refused. +- Statistics are not used under a row filter, which also turns off `LIMIT` + pushdown, so such queries read the data. +- Time travel, branches, system tables, incremental and audit-log reads, and + search are refused on such tables. + ## Python Multimodal Helper Functions When you use `pypaimon_rust.datafusion.SQLContext`, the Python binding registers a small set of scalar helper functions for BLOB-backed media and vector workflows. These helpers are Python-binding built-ins; they are not registered by the Rust `paimon_datafusion::SQLContext`.