From 91e5e94611154e5e3fc088807528e4670d7ebc84 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Sun, 26 Jul 2026 07:50:57 -0400 Subject: [PATCH] [table] Support query-auth row filters and column masking at read time --- Cargo.lock | 1 + crates/integrations/datafusion/Cargo.toml | 2 + .../integrations/datafusion/src/merge_into.rs | 27 +- .../datafusion/src/physical_plan/scan.rs | 24 + .../integrations/datafusion/src/table/mod.rs | 4 + .../datafusion/src/variant_pushdown.rs | 1 + .../datafusion/tests/query_auth.rs | 212 ++++ crates/paimon/src/arrow/residual.rs | 4 +- crates/paimon/src/spec/schema.rs | 13 + .../table/btree_global_index_build_builder.rs | 10 +- .../paimon/src/table/bucket_assigner_cross.rs | 7 + crates/paimon/src/table/cow_writer.rs | 12 +- .../paimon/src/table/data_evolution_writer.rs | 15 +- .../paimon/src/table/format_read_builder.rs | 30 +- crates/paimon/src/table/format_table_read.rs | 9 +- crates/paimon/src/table/format_table_scan.rs | 87 +- .../src/table/full_text_search_builder.rs | 10 +- .../paimon/src/table/hybrid_search_builder.rs | 9 +- .../src/table/lumina_index_build_builder.rs | 7 +- crates/paimon/src/table/mod.rs | 190 ++- crates/paimon/src/table/partition_stat.rs | 5 + crates/paimon/src/table/query_auth.rs | 1085 +++++++++++++++++ crates/paimon/src/table/read_builder.rs | 264 +++- crates/paimon/src/table/rest_env.rs | 9 + crates/paimon/src/table/source.rs | 164 ++- crates/paimon/src/table/table_commit.rs | 10 + crates/paimon/src/table/table_read.rs | 357 +++++- crates/paimon/src/table/table_scan.rs | 186 ++- .../paimon/src/table/vector_search_builder.rs | 38 +- .../src/table/vindex_index_build_builder.rs | 7 +- crates/paimon/tests/mock_server.rs | 31 +- crates/paimon/tests/rest_catalog_test.rs | 560 ++++++++- 32 files changed, 3285 insertions(+), 105 deletions(-) create mode 100644 crates/integrations/datafusion/tests/query_auth.rs create mode 100644 crates/paimon/src/table/query_auth.rs diff --git a/Cargo.lock b/Cargo.lock index 7ee0257e..4ad09ac8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4606,6 +4606,7 @@ dependencies = [ "arrow-schema", "arrow-select", "async-trait", + "axum", "bytes", "chrono", "constant_time_eq", diff --git a/crates/integrations/datafusion/Cargo.toml b/crates/integrations/datafusion/Cargo.toml index 17acaea0..27e34402 100644 --- a/crates/integrations/datafusion/Cargo.toml +++ b/crates/integrations/datafusion/Cargo.toml @@ -50,6 +50,8 @@ uuid = { version = "1", features = ["v4"] } [dev-dependencies] arrow-array = { workspace = true } arrow-schema = { workspace = true } +# for the shared REST catalog mock (crates/paimon/tests/mock_server.rs) +axum = { version = "0.7", features = ["macros", "tokio", "http1", "http2"] } bytes = "1.7.1" flate2 = "1" paimon-ftindex-core = "0.1.0" diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index 97c36079..1d3e4f8c 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -1233,12 +1233,14 @@ pub(crate) async fn register_cow_target_table( return Ok((false, table_name)); } - // Read all files in parallel - let read_futures: Vec<_> = file_index - .iter() - .enumerate() - .map(|(file_idx, file_info)| async move { - let single_split = DataSplitBuilder::new() + // This read rewrites the target files, so it must see raw rows: authorize an + // unrestricted grant once and stamp it on every split. Stamping the scan + // plan's grant instead would shift the positional row offsets the writer + // replays under a row filter, rewriting the wrong rows. + let mut splits = Vec::with_capacity(file_index.len()); + for file_info in file_index.iter() { + splits.push( + DataSplitBuilder::new() .with_snapshot(file_info.snapshot_id) .with_partition( paimon::spec::BinaryRow::from_serialized_bytes(&file_info.partition) @@ -1249,8 +1251,19 @@ pub(crate) async fn register_cow_target_table( .with_total_buckets(file_info.total_buckets) .with_data_files(vec![file_info.file_meta.clone()]) .build() - .map_err(to_datafusion_error)?; + .map_err(to_datafusion_error)?, + ); + } + let splits = table + .authorize_rewrite_splits(splits) + .await + .map_err(to_datafusion_error)?; + // Read all files in parallel + let read_futures: Vec<_> = splits + .into_iter() + .enumerate() + .map(|(file_idx, single_split)| async move { let read = table .new_read_builder() .new_read() diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index bf8373dd..f713e572 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -851,6 +851,18 @@ impl PaimonTableScan { return Statistics::unknown_column(&self.schema()); } + // Manifest stats describe the data BEFORE query-auth enforcement: the + // bounds of a masked column are the raw values the mask hides, and null + // counts cover rows the row filter drops. Publishing them (e.g. via + // EXPLAIN, or to the optimizer) would leak exactly what enforcement hides. + if partitions + .iter() + .flat_map(|splits| splits.iter()) + .any(DataSplit::has_restricted_query_auth_grant) + { + return Statistics::unknown_column(&self.schema()); + } + let Ok(merge_engine) = self.table.schema().core_options().merge_engine() else { return Statistics::unknown_column(&self.schema()); }; @@ -1066,6 +1078,18 @@ impl ExecutionPlan for PaimonTableScan { None => &self.planned_partitions, }; + // Under a restricted grant the manifest row count is the count BEFORE + // the row filter runs, so publishing it (EXPLAIN, the optimizer) would + // disclose how much data the filter hides — the same reason the column + // statistics are suppressed. Report the whole thing as unknown. + if partitions + .iter() + .flat_map(|splits| splits.iter()) + .any(DataSplit::has_restricted_query_auth_grant) + { + return Ok(Arc::new(Statistics::new_unknown(&self.schema()))); + } + let mut total_rows: usize = 0; let mut all_row_counts_known = true; for splits in partitions { diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 1bfc47fc..618f7f85 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -450,7 +450,11 @@ impl TableProvider for PaimonTableProvider { .map_err(to_datafusion_error)?; let target = state.config_options().execution.target_partitions; + // Inexact plan row counts (a query-auth row filter drops rows inside + // `TableRead`) would let DataFusion's aggregate-statistics rule answer + // COUNT(*) with the unfiltered count without ever invoking the read. let filter_exact = !filter_analysis.requires_residual + && plan.row_counts_exact() && filter_analysis .pushed_predicate .as_ref() diff --git a/crates/integrations/datafusion/src/variant_pushdown.rs b/crates/integrations/datafusion/src/variant_pushdown.rs index a5bb3482..1c64daaa 100644 --- a/crates/integrations/datafusion/src/variant_pushdown.rs +++ b/crates/integrations/datafusion/src/variant_pushdown.rs @@ -243,6 +243,7 @@ impl ExtensionPlanner for VariantExtractionExtensionPlanner { .collect() }; let filter_exact = !filter_analysis.requires_residual + && plan.row_counts_exact() && filter_analysis .pushed_predicate .as_ref() diff --git a/crates/integrations/datafusion/tests/query_auth.rs b/crates/integrations/datafusion/tests/query_auth.rs new file mode 100644 index 00000000..46250d69 --- /dev/null +++ b/crates/integrations/datafusion/tests/query_auth.rs @@ -0,0 +1,212 @@ +// 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. + +//! Query-auth enforcement through the DataFusion provider: SQL over a REST +//! catalog table with `query-auth.enabled` must apply the per-user grant +//! (row filter + column masking) fetched at scan-plan time, and COUNT(*) +//! must not shortcut to unfiltered statistics. +//! +//! This is the same provider path the Python binding +//! (`pypaimon_rust.datafusion`) drives via FFI. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Int32Array, Int64Array, StringArray}; +use datafusion::arrow::compute::cast; +use datafusion::arrow::datatypes::{ + DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, +}; +use datafusion::arrow::record_batch::RecordBatch; +use paimon::api::{AuthTableQueryResponse, ConfigResponse}; +use paimon::catalog::{Identifier, RESTCatalog}; +use paimon::spec::{BigIntType, DataType, IntType, Schema, VarCharType}; +use paimon::{Catalog, CatalogOptions, FileSystemCatalog, Options, Table}; +use paimon_datafusion::SQLContext; + +// Shared REST catalog mock (same source the paimon crate's integration tests +// use); only a subset of its helpers is exercised here. +#[allow(dead_code)] +#[path = "../../../paimon/tests/mock_server.rs"] +mod mock_server; +use mock_server::start_mock_server; + +async fn write_batch(table: &Table, batch: RecordBatch, commit_user: &str) { + let write_builder = table + .new_write_builder() + .with_commit_user(commit_user) + .expect("valid commit user"); + let mut write = write_builder.new_write().expect("create writer"); + write.write_arrow_batch(&batch).await.expect("write batch"); + let messages = write.prepare_commit().await.expect("prepare commit"); + write_builder + .new_commit() + .commit(messages) + .await + .expect("commit batch"); +} + +// Multi-threaded: the provider's `block_on_with_runtime` bridges park the +// current thread, which must not be the only thread serving the mock server. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_query_auth_grant_enforced_via_sql() { + let tmp = tempfile::tempdir().unwrap(); + let warehouse = format!("file://{}", tmp.path().display()); + + // Write demo_employees (id, name, salary) through a plain FileSystemCatalog. + let mut fs_options = Options::new(); + fs_options.set(CatalogOptions::WAREHOUSE, &warehouse); + let fs_catalog = FileSystemCatalog::new(fs_options).expect("create filesystem catalog"); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let columns = || { + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("salary", DataType::BigInt(BigIntType::new())) + .option("bucket", "1") + .option("bucket-key", "id") + }; + let identifier = Identifier::new("default", "qa_emp"); + fs_catalog + .create_table(&identifier, columns().build().unwrap(), false) + .await + .unwrap(); + let writer = fs_catalog.get_table(&identifier).await.unwrap(); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("salary", ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec![ + "alice", "bob", "charlie", "diana", "eve", + ])), + Arc::new(Int64Array::from(vec![120000, 85000, 95000, 70000, 99000])), + ], + ) + .unwrap(); + write_batch(&writer, batch, "u1").await; + + // Serve the same files through a mock REST catalog whose grant applies a + // row filter (salary >= 90000) and masks name -> UPPER(name). + let mut defaults = HashMap::new(); + defaults.insert("prefix".to_string(), "mock-test".to_string()); + let server = start_mock_server( + "test_warehouse".to_string(), + "/tmp/test_warehouse".to_string(), + ConfigResponse::new(defaults), + vec!["default".to_string()], + ) + .await; + server.add_table_with_schema( + "default", + "qa_emp", + columns() + .option("query-auth.enabled", "true") + .build() + .unwrap(), + &format!("{warehouse}/default.db/qa_emp"), + ); + server.set_auth_response( + "default", + "qa_emp", + AuthTableQueryResponse { + filter: Some(vec![ + r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":2,"name":"salary","type":"BIGINT"}},"function":"GREATER_OR_EQUAL","literals":[90000]}"# + .to_string(), + ]), + column_masking: Some(HashMap::from([( + "name".to_string(), + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + .to_string(), + )])), + }, + ); + + let mut rest_options = Options::new(); + rest_options.set("uri", server.url().expect("server url")); + rest_options.set("warehouse", "test_warehouse"); + rest_options.set("token.provider", "bear"); + rest_options.set("token", "test_token"); + let rest_catalog = RESTCatalog::new(rest_options, true) + .await + .expect("create REST catalog"); + + // Sanity: the mock-served table must resolve through the Catalog trait. + rest_catalog + .get_table(&identifier) + .await + .expect("REST get_table(default.qa_emp)"); + + let mut ctx = SQLContext::new(); + ctx.register_catalog("paimon", Arc::new(rest_catalog)) + .await + .expect("register catalog"); + + // Row filter + masking must both be applied on the SQL result. + let batches = ctx + .sql("SELECT id, name, salary FROM paimon.default.qa_emp ORDER BY id") + .await + .expect("plan select") + .collect() + .await + .expect("execute select"); + let mut rows = Vec::new(); + for b in &batches { + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + // DataFusion may hand back Utf8View for string columns; normalize. + let names = cast(b.column(1), &ArrowDataType::Utf8).expect("cast name to Utf8"); + let names = names.as_any().downcast_ref::().unwrap(); + let sal = b.column(2).as_any().downcast_ref::().unwrap(); + for r in 0..b.num_rows() { + rows.push((ids.value(r), names.value(r).to_string(), sal.value(r))); + } + } + assert_eq!( + rows, + vec![ + (1, "ALICE".to_string(), 120000), + (3, "CHARLIE".to_string(), 95000), + (5, "EVE".to_string(), 99000), + ], + "grant must drop salary<90000 rows and uppercase name" + ); + + // COUNT(*) must reflect the filtered row count, not unfiltered statistics + // (the aggregate-statistics optimization must be disabled by the inexact + // plan row counts). + let batches = ctx + .sql("SELECT COUNT(*) FROM paimon.default.qa_emp") + .await + .expect("plan count") + .collect() + .await + .expect("execute count"); + let count = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(count, 3, "COUNT(*) must not use unfiltered statistics"); +} diff --git a/crates/paimon/src/arrow/residual.rs b/crates/paimon/src/arrow/residual.rs index 60cfe401..edc6d1fc 100644 --- a/crates/paimon/src/arrow/residual.rs +++ b/crates/paimon/src/arrow/residual.rs @@ -566,7 +566,7 @@ fn evaluate_set_membership_predicate( Ok(combined) } -fn evaluate_column_predicate( +pub(crate) fn evaluate_column_predicate( column: &ArrayRef, scalar: &Scalar, op: PredicateOperator, @@ -725,7 +725,7 @@ fn combine_filter_masks(left: &BooleanArray, right: &BooleanArray, use_or: bool) BooleanArray::new(values, None) } -fn boolean_mask_from_predicate( +pub(crate) fn boolean_mask_from_predicate( len: usize, mut predicate: impl FnMut(usize) -> bool, ) -> BooleanArray { diff --git a/crates/paimon/src/spec/schema.rs b/crates/paimon/src/spec/schema.rs index 0245cf41..b5f7b565 100644 --- a/crates/paimon/src/spec/schema.rs +++ b/crates/paimon/src/spec/schema.rs @@ -164,6 +164,19 @@ impl TableSchema { new_schema } + /// Force `query-auth.enabled = true` on this schema. + /// + /// Used when a copy must inherit the flag from a table whose schema carries + /// it (the REST catalog delivers it on the table response, so an on-disk + /// schema — e.g. a branch's — need not have it). + pub(crate) fn copy_with_query_auth_enabled(&self) -> Self { + let mut new_schema = self.clone(); + new_schema + .options + .insert(QUERY_AUTH_ENABLED_OPTION.to_string(), "true".to_string()); + new_schema + } + /// Apply a list of schema changes and return a new schema with incremented ID. /// /// Column-level changes operate on **top-level** columns only: a diff --git a/crates/paimon/src/table/btree_global_index_build_builder.rs b/crates/paimon/src/table/btree_global_index_build_builder.rs index 693a15dc..3278c8c4 100644 --- a/crates/paimon/src/table/btree_global_index_build_builder.rs +++ b/crates/paimon/src/table/btree_global_index_build_builder.rs @@ -614,7 +614,15 @@ async fn extract_index_rows( index_field: &DataField, serialize_key: SerializeKeyFn, ) -> Result> { - let splits = build_read_splits_for_shard(shard)?; + // Building the global index reads the indexed column across the shard. Under + // a restricted query-auth grant that read would drop/mask rows, so the index + // would be built over a filtered view. Require an unrestricted grant and read + // raw (stamp the returned grant on each split). + let write_grant = table.authorize_unrestricted_read().await?; + let splits: Vec = build_read_splits_for_shard(shard)? + .into_iter() + .map(|s| s.with_query_auth_grant(write_grant.clone())) + .collect(); let mut read_builder = table.new_read_builder(); read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; diff --git a/crates/paimon/src/table/bucket_assigner_cross.rs b/crates/paimon/src/table/bucket_assigner_cross.rs index f0e711b2..83f472ea 100644 --- a/crates/paimon/src/table/bucket_assigner_cross.rs +++ b/crates/paimon/src/table/bucket_assigner_cross.rs @@ -85,6 +85,13 @@ impl GlobalPartitionIndex { .collect(); let projected_pk_indices: Vec = (0..pk_fields.len()).collect(); + // Building the cross-partition PK index reads every primary key. Under a + // restricted query-auth grant the planned splits would filter rows, so + // the index would miss hidden keys and produce duplicate PKs / lost + // updates on upsert. Require a fully unrestricted grant (the plan then + // stamps that unrestricted grant on its splits, so the reads run raw). + table.authorize_unrestricted_read().await?; + let mut rb = table.new_read_builder(); rb.with_projection(&pk_field_names)?; let scan = rb.new_scan().with_scan_all_files(); diff --git a/crates/paimon/src/table/cow_writer.rs b/crates/paimon/src/table/cow_writer.rs index 6c348d15..6bc6f9bc 100644 --- a/crates/paimon/src/table/cow_writer.rs +++ b/crates/paimon/src/table/cow_writer.rs @@ -210,6 +210,14 @@ impl CopyOnWriteMergeWriter { return Ok(Vec::new()); } + // A copy-on-write rewrite reads each affected file and rewrites it in + // place. Under a restricted query-auth grant that read would filter and + // mask rows, silently deleting hidden rows and persisting masked values + // on commit — so require a fully unrestricted grant and read raw. The + // returned grant (unrestricted, or `None` for a non-query-auth table) is + // stamped on each split so `to_arrow` reads raw instead of failing closed. + let write_grant = self.table.authorize_unrestricted_read().await?; + let schema = self.table.schema(); let core_options = CoreOptions::new(schema.options()); let partition_keys: Vec = schema.partition_keys().to_vec(); @@ -232,6 +240,7 @@ impl CopyOnWriteMergeWriter { let update_batches = &self.update_batches; let file_index = &self.file_index; let table = &self.table; + let write_grant = &write_grant; let partition_keys = &partition_keys; let partition_computer = &partition_computer; let write_fields = &write_fields; @@ -253,7 +262,8 @@ impl CopyOnWriteMergeWriter { .with_bucket_path(file_info.bucket_path.clone()) .with_total_buckets(file_info.total_buckets) .with_data_files(vec![file_info.file_meta.clone()]) - .build()?; + .build()? + .with_query_auth_grant(write_grant.clone()); let read = table.new_read_builder().new_read()?; let original_batches: Vec = diff --git a/crates/paimon/src/table/data_evolution_writer.rs b/crates/paimon/src/table/data_evolution_writer.rs index 47752e3f..7b7b7701 100644 --- a/crates/paimon/src/table/data_evolution_writer.rs +++ b/crates/paimon/src/table/data_evolution_writer.rs @@ -164,6 +164,12 @@ impl DataEvolutionWriter { return Ok(Vec::new()); } + // This rewrite reads each affected file's original columns and rewrites + // it; under a restricted query-auth grant that read would filter/mask + // rows into the committed result. Require an unrestricted grant and read + // raw (stamp the returned grant on each split). + let write_grant = self.table.authorize_unrestricted_read().await?; + // 1. Scan file metadata and build row_id -> file group index. // In data-evolution tables, multiple files can share the same first_row_id // (base file + partial-column files). We must group them so the reader @@ -248,7 +254,8 @@ impl DataEvolutionWriter { .with_total_buckets(file_range.total_buckets) .with_data_files(file_range.files.clone()) .with_raw_convertible(file_range.files.len() == 1) - .build()?; + .build()? + .with_query_auth_grant(write_grant.clone()); let stream = read.to_arrow(&[split])?; let original_batches: Vec = stream.try_collect().await?; @@ -460,6 +467,12 @@ impl DataEvolutionDeleteWriter { return Ok(Vec::new()); } + // The row ids to delete come from a read the caller performed; under a + // restricted grant that read hid rows, so the committed deletion vectors + // would encode a delete set derived from a filtered view. Require an + // unrestricted grant, like `DataEvolutionWriter::prepare_commit`. + self.table.authorize_unrestricted_read().await?; + let scan = self .table .new_read_builder() diff --git a/crates/paimon/src/table/format_read_builder.rs b/crates/paimon/src/table/format_read_builder.rs index f97805c0..2e99f6a8 100644 --- a/crates/paimon/src/table/format_read_builder.rs +++ b/crates/paimon/src/table/format_read_builder.rs @@ -24,6 +24,7 @@ use super::{Table, TableRead, TableScan}; use crate::spec::{DataField, Predicate}; use crate::table::source::RowRange; use crate::Result; +use std::collections::HashSet; #[derive(Debug, Clone)] pub(crate) struct FormatReadBuilder<'a> { @@ -37,6 +38,7 @@ pub(crate) struct FormatReadBuilder<'a> { data_predicates: Vec, limit: Option, case_sensitive: bool, + filter_columns: HashSet, } impl<'a> FormatReadBuilder<'a> { @@ -49,6 +51,7 @@ impl<'a> FormatReadBuilder<'a> { data_predicates: Vec::new(), limit: None, case_sensitive: true, + filter_columns: HashSet::new(), } } @@ -81,6 +84,10 @@ impl<'a> FormatReadBuilder<'a> { } pub(crate) fn with_filter(&mut self, filter: Predicate) -> &mut Self { + // Capture the full predicate's columns before it is split, so masked and + // out-of-scope partition keys can't prune on their raw value. + self.filter_columns.clear(); + filter.collect_leaf_field_indices(&mut self.filter_columns); let (partition_predicate, data_predicates) = split_scan_predicates(self.table, filter); self.partition_filter = partition_predicate.map(|pred| { PartitionFilter::from_predicate(pred, &self.table.schema().partition_fields()) @@ -111,11 +118,30 @@ impl<'a> FormatReadBuilder<'a> { self.limit, None, ) + .with_query_auth_scope(self.filter_columns.clone(), self.projected_schema_indices()) + } + + /// Table-schema indices of the projected columns (`None` = all). + fn projected_schema_indices(&self) -> Option> { + // Resolve names too (see `PaimonReadBuilder::projected_schema_indices`). + self.resolve_read_type().ok().flatten().map(|fields| { + fields + .iter() + .filter_map(|f| { + self.table + .schema() + .fields() + .iter() + .position(|s| s.id() == f.id()) + }) + .collect() + }) } pub(crate) fn new_read(&self) -> Result> { - let core_options = self.table.schema().core_options(); - core_options.ensure_read_authorized()?; + // Query-auth is enforced in `TableRead::to_arrow` off the grant stamped + // on the splits by planning; no gate needed here (see the Paimon + // `PaimonReadBuilder::new_read`). let read_type = match self.resolve_read_type()? { None => self.table.schema().fields().to_vec(), Some(fields) => fields, diff --git a/crates/paimon/src/table/format_table_read.rs b/crates/paimon/src/table/format_table_read.rs index b9ea8e5c..17ac09d9 100644 --- a/crates/paimon/src/table/format_table_read.rs +++ b/crates/paimon/src/table/format_table_read.rs @@ -70,6 +70,10 @@ impl<'a> FormatTableRead<'a> { self.table } + pub(crate) fn limit(&self) -> Option { + self.limit + } + pub(crate) fn with_filter(mut self, filter: Predicate) -> Self { self.data_predicates = split_scan_predicates(self.table, filter).1; self @@ -87,8 +91,9 @@ impl<'a> FormatTableRead<'a> { &self, data_splits: &[DataSplit], ) -> crate::Result { - let core_options = self.table.schema().core_options(); - core_options.ensure_read_authorized()?; + // Query-auth (fail-closed + row filter + masking) is enforced by the + // outer `TableRead::to_arrow` off the grant stamped on the splits. + let core_options = self.table.schema.core_options(); let read_type = self.read_type.clone(); let output_schema = build_target_arrow_schema(&read_type)?; let partition_keys = self.table.schema().partition_keys().to_vec(); diff --git a/crates/paimon/src/table/format_table_scan.rs b/crates/paimon/src/table/format_table_scan.rs index 28907664..6d1ba52e 100644 --- a/crates/paimon/src/table/format_table_scan.rs +++ b/crates/paimon/src/table/format_table_scan.rs @@ -32,6 +32,8 @@ pub(crate) struct FormatTableScan<'a> { table: &'a Table, partition_filter: Option, limit: Option, + query_auth_filter_columns: std::collections::HashSet, + query_auth_projected: Option>, } impl<'a> FormatTableScan<'a> { @@ -44,26 +46,87 @@ impl<'a> FormatTableScan<'a> { table, partition_filter, limit, + query_auth_filter_columns: std::collections::HashSet::new(), + query_auth_projected: None, } } + pub(super) fn with_query_auth_scope( + mut self, + filter_columns: std::collections::HashSet, + projected: Option>, + ) -> Self { + self.query_auth_filter_columns = filter_columns; + self.query_auth_projected = projected; + self + } + pub(crate) async fn plan(&self) -> crate::Result { - self.ensure_query_auth_allowed()?; - self.plan_inner(None).await + let grant = self.ensure_query_auth_allowed().await?; + let has_row_filter = grant.as_deref().is_some_and(|g| g.has_row_filter()); + self.plan_inner(None, has_row_filter) + .await + .map(|plan| self.finalize_plan(plan, grant.as_ref())) } pub(crate) async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { - self.ensure_query_auth_allowed()?; + let grant = self.ensure_query_auth_allowed().await?; + let has_row_filter = grant.as_deref().is_some_and(|g| g.has_row_filter()); let mut trace = ScanTrace::default(); - let plan = self.plan_inner(Some(&mut trace)).await?; - Ok((plan, trace)) + let plan = self.plan_inner(Some(&mut trace), has_row_filter).await?; + Ok((self.finalize_plan(plan, grant.as_ref()), trace)) } - fn ensure_query_auth_allowed(&self) -> crate::Result<()> { - CoreOptions::new(self.table.schema().options()).ensure_read_authorized() + /// Stamp the grant onto every split (so `TableRead::to_arrow` enforces it) + /// and mark row counts inexact when it carries a row filter. + fn finalize_plan( + &self, + plan: Plan, + grant: Option<&std::sync::Arc>, + ) -> Plan { + // Any restricted grant invalidates the plan's statistics (see the + // Paimon `TableScan::finalize_plan`). + let restricted = grant.is_some_and(|g| g.has_server_restrictions()); + let plan = plan.stamp_query_auth_grant(grant.cloned()); + if restricted { + plan.with_inexact_row_counts() + } else { + plan + } } - async fn plan_inner(&self, trace: Option<&mut ScanTrace>) -> crate::Result { + async fn ensure_query_auth_allowed( + &self, + ) -> crate::Result>> { + // Fetch/refresh the grant at plan time (Java parity), then guard + // against pruning on masked or out-of-scope columns. + let select = self.query_auth_projected.as_ref().map(|projected| { + projected + .iter() + .copied() + .chain(self.query_auth_filter_columns.iter().copied()) + .collect::>() + }); + let grant = self + .table + .verify_query_auth_for_read(select.as_ref()) + .await?; + if let Some(grant) = &grant { + crate::table::query_auth::scope_check( + grant, + self.table.schema().fields(), + &self.query_auth_filter_columns, + self.query_auth_projected.clone(), + )?; + } + Ok(grant) + } + + async fn plan_inner( + &self, + trace: Option<&mut ScanTrace>, + query_auth_row_filter: bool, + ) -> crate::Result { let core_options = CoreOptions::new(self.table.schema().options()); let format_extension = supported_format_table_extension(core_options.file_format())?; let schema_id = self.table.schema().id(); @@ -103,7 +166,7 @@ impl<'a> FormatTableScan<'a> { .cmp(&right.data_files()[0].file_name) }) }); - splits = self.apply_limit_pushdown(splits); + splits = self.apply_limit_pushdown(splits, query_auth_row_filter); if let Some(trace) = trace { trace.record_final_plan(splits.len(), splits.len(), splits.len()); @@ -275,7 +338,13 @@ impl<'a> FormatTableScan<'a> { pub(crate) fn apply_limit_pushdown( &self, splits: Vec, + query_auth_row_filter: bool, ) -> Vec { + // A query-auth row filter runs as a residual pass at read time, so the + // scan must not cap files by an unfiltered limit before that. + if query_auth_row_filter { + return splits; + } match self.limit { Some(0) => Vec::new(), Some(limit) if splits.len() > limit => splits.into_iter().take(limit).collect(), diff --git a/crates/paimon/src/table/full_text_search_builder.rs b/crates/paimon/src/table/full_text_search_builder.rs index 702c3295..4f8d78ba 100644 --- a/crates/paimon/src/table/full_text_search_builder.rs +++ b/crates/paimon/src/table/full_text_search_builder.rs @@ -116,8 +116,10 @@ impl<'a> FullTextSearchBuilder<'a> { pub async fn execute_scored(&self) -> crate::Result { // Fail closed: returns data-derived row ranges outside `TableScan`/`TableRead`. + // Strict: search results bypass the query-auth row filter, so only a + // fully unrestricted grant may search. + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let text_column = self.text_column .as_deref() @@ -202,9 +204,11 @@ impl<'a> FullTextSearchBuilder<'a> { /// returning nothing, since the append/data-evolution materialized read is not /// supported here. pub async fn execute_read(&self) -> crate::Result { - // Fail closed: returns data outside `TableScan`/`TableRead`. + // Fail closed: materializes rows outside `TableScan`/`TableRead`, so it + // cannot apply the row filter / masking — only a fully unrestricted + // grant may run it (same gate as the scored entry points). + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let text_column = self.text_column .as_deref() diff --git a/crates/paimon/src/table/hybrid_search_builder.rs b/crates/paimon/src/table/hybrid_search_builder.rs index ab290820..f80e5beb 100644 --- a/crates/paimon/src/table/hybrid_search_builder.rs +++ b/crates/paimon/src/table/hybrid_search_builder.rs @@ -286,8 +286,10 @@ impl<'a> HybridSearchBuilder<'a> { } pub async fn execute_scored(&self) -> crate::Result { + // Strict: search results bypass the query-auth row filter, so only a + // fully unrestricted grant may search. + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { message: "Limit must be set via with_limit()".to_string(), })?; @@ -350,8 +352,11 @@ impl<'a> HybridSearchBuilder<'a> { /// (append/data-evolution) hybrid is unsupported here — those use /// `execute`/`execute_scored`. Mirrors Java `HybridSearchBuilderImpl` PK path. pub async fn execute_read(&self) -> crate::Result { + // Materializes rows outside `TableScan`/`TableRead`, so it cannot apply + // the row filter / masking — only a fully unrestricted grant may run it + // (same gate as the vector and full-text builders). + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let limit = self.limit.ok_or_else(|| crate::Error::ConfigInvalid { message: "Limit must be set via with_limit()".to_string(), })?; diff --git a/crates/paimon/src/table/lumina_index_build_builder.rs b/crates/paimon/src/table/lumina_index_build_builder.rs index 5e45abf8..69502d59 100644 --- a/crates/paimon/src/table/lumina_index_build_builder.rs +++ b/crates/paimon/src/table/lumina_index_build_builder.rs @@ -571,6 +571,10 @@ async fn extract_vectors( index_column: &str, dimension: i32, ) -> Result> { + // Index building reads raw values, so it requires an unrestricted grant (a + // restricted one would index a filtered/masked view); stamp it so the read + // is authorized. Mirrors the B-tree index builder. + let build_grant = table.authorize_unrestricted_read().await?; let split = DataSplitBuilder::new() .with_snapshot(shard.snapshot_id) .with_partition(shard.partition.clone()) @@ -582,7 +586,8 @@ async fn extract_vectors( shard.row_range_start, shard.row_range_end, )]) - .build()?; + .build()? + .with_query_auth_grant(build_grant); let mut read_builder = table.new_read_builder(); read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index be680cb0..9e653990 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -74,6 +74,7 @@ mod pk_vector_position_read; mod pk_vector_scan; mod postpone_file_writer; mod prepared_files; +pub(crate) mod query_auth; mod read_builder; pub mod referenced_files; pub(crate) mod rest_env; @@ -142,7 +143,9 @@ pub use write_builder::WriteBuilder; use crate::catalog::{validate_branch_name, Identifier, DEFAULT_MAIN_BRANCH}; use crate::io::FileIO; use crate::spec::{CoreOptions, DataField, Snapshot, TableSchema}; -use std::collections::HashMap; +use query_auth::QueryAuthGrant; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; /// Table represents a table in the catalog. #[derive(Debug, Clone)] @@ -189,6 +192,130 @@ impl Table { } } + /// Authorize this user against the REST server when `query-auth.enabled` is + /// set: fetch and parse the per-user row filter / column masking and return + /// it as the grant the read pipeline enforces. `select` are the queried + /// columns (table-schema indices; `None` = all) — like Java's + /// `readType.getFieldNames()`, so a column-restricted user can still read + /// an authorized subset; the grant is scoped to exactly those columns and a + /// wider read fails closed until it re-plans. Returns `None` when the table + /// is not `query-auth.enabled`. Called from scan planning and search + /// execution at the same per-plan frequency as Java's + /// `CatalogEnvironment.tableQueryAuth()`, so a revoked grant takes effect on + /// the next plan. The grant is threaded to the read on the split it plans + /// (never a shared mutable slot), so it is per-query and cannot leak into a + /// concurrent query or a write-path rewrite. + pub(crate) async fn verify_query_auth_for_read( + &self, + select: Option<&HashSet>, + ) -> Result>> { + if !CoreOptions::new(self.schema.options()).query_auth_enabled() { + return Ok(None); + } + let Some(rest_env) = &self.rest_env else { + return Err(crate::Error::Unsupported { + message: "reading a table with 'query-auth.enabled' = true requires a REST \ + catalog to authorize the query" + .to_string(), + }); + }; + let fields = self.schema.fields(); + let select_names = select.map(|indices| { + fields + .iter() + .enumerate() + .filter(|(i, _)| indices.contains(i)) + .map(|(_, f)| f.name().to_string()) + .collect::>() + }); + let auth = rest_env.table_query_auth(select_names).await?; + let filters = query_auth::parse_auth_filters(&auth.filter.unwrap_or_default(), fields)?; + let masks = + query_auth::parse_column_masking(&auth.column_masking.unwrap_or_default(), fields)?; + let grant = + QueryAuthGrant::new(filters, masks, select.cloned()).with_schema_id(self.schema.id()); + + // The server expresses its rules against the table's CURRENT schema, but + // a time-travelled or branch copy reads a DIFFERENT one. Binding those + // rules by name onto it is unsound — a column dropped and re-added under + // the same name has a different field id, so a filter or mask would bind + // to an unrelated historical column. Fail closed rather than mis-bind. + if !grant.is_unrestricted() && (self.time_traveled || self.branch_reference) { + return Err(crate::Error::Unsupported { + message: "a query-auth row filter / column masking grant cannot be applied to a \ + time-travelled or branch read: the grant is bound to the table's \ + current schema" + .to_string(), + }); + } + + Ok(Some(Arc::new(grant))) + } + + /// Authorize a read that cannot enforce a row filter / masking on its output + /// (search, system tables, or a write-path rewrite that must read raw). + /// Returns the grant to stamp on the splits it reads (`None` = not a + /// query-auth table). Fails closed when the grant is restricted — such a + /// path must never run under a partial grant, since it would either leak + /// (search/metadata) or silently drop/mask rows into a committed rewrite. + pub(crate) async fn authorize_unrestricted_read(&self) -> Result>> { + match self.verify_query_auth_for_read(None).await? { + Some(grant) if grant.is_unrestricted() => Ok(Some(grant)), + Some(_) => { + // A restricted grant only arrives on a query-auth.enabled table, + // where the strict schema check always fails closed. + CoreOptions::new(self.schema.options()).ensure_read_authorized()?; + Ok(None) + } + None => Ok(None), + } + } + + /// Authorize an internal read that rewrites data (copy-on-write DML, index + /// builds) and stamp the grant on `splits`. + /// + /// Such a read must see raw rows — rewriting from a filtered or masked view + /// would destroy hidden rows and persist masked values — so it requires a + /// fully unrestricted grant and fails closed otherwise. The grant the caller + /// would get from scan planning must NOT be used here: under a row filter it + /// shifts row offsets, and rewrites replay those offsets. + pub async fn authorize_rewrite_splits(&self, splits: Vec) -> Result> { + // This is the only public API that stamps a grant, so it must not be + // usable to launder splits: refuse ones that already carry a restricted + // grant rather than overwriting it. + if splits + .iter() + .any(DataSplit::has_restricted_query_auth_grant) + { + return Err(crate::Error::Unsupported { + message: "cannot re-authorize a split already planned under a query-auth row \ + filter / column masking grant" + .to_string(), + }); + } + let grant = self.authorize_unrestricted_read().await?; + Ok(splits + .into_iter() + .map(|split| split.with_query_auth_grant(grant.clone())) + .collect()) + } + + /// Authorize a commit. Writes must require a fully unrestricted grant: the + /// data being committed may have come from an enforced read (e.g. `INSERT + /// OVERWRITE t SELECT * FROM t`), which would destroy rows the grant hid and + /// persist masked values as the stored data. Committing cannot tell where + /// its rows came from, so it fails closed for any restricted grant. + pub(crate) async fn authorize_unrestricted_write(&self) -> Result<()> { + self.authorize_unrestricted_read().await.map(|_| ()) + } + + /// Fail closed when a read reaches `TableRead::to_arrow` without a grant + /// stamped on its splits (an unauthorized path) and the table is + /// `query-auth.enabled`; a no-op otherwise. + pub(crate) fn ensure_read_without_grant(&self) -> Result<()> { + CoreOptions::new(self.schema.options()).ensure_read_authorized() + } + /// Get the table's identifier. pub fn identifier(&self) -> &Identifier { &self.identifier @@ -416,11 +543,22 @@ impl Table { })?; let mut options = schema.options().clone(); options.insert("branch".to_string(), branch.clone()); + // `query-auth.enabled` is delivered by the REST catalog on the table + // response, so the branch's on-disk schema need not carry it. Inherit it + // from this table: a branch references the same data files, and dropping + // the flag here would make `t$branch_x` read them raw, unauthorized. + let branch_schema = if CoreOptions::new(self.schema.options()).query_auth_enabled() { + schema + .copy_with_replaced_options(options) + .copy_with_query_auth_enabled() + } else { + schema.copy_with_replaced_options(options) + }; Ok(Self { file_io: self.file_io.clone(), identifier: self.identifier.clone(), location: self.location.clone(), - schema: schema.copy_with_replaced_options(options), + schema: branch_schema, schema_manager, branch, branch_reference: true, @@ -482,3 +620,51 @@ pub(crate) fn query_auth_table() -> Table { None, ) } + +#[cfg(test)] +mod tests { + #[tokio::test] + async fn test_authorize_unrestricted_read_fails_closed() { + // Every write/build path (cow_writer, data_evolution_writer, btree index + // build, cross-partition bucket assign) authorizes through this gate + // before its internal read. A query-auth table it cannot authorize as + // unrestricted must fail closed rather than read raw and rewrite + // filtered/masked data into a committed result. + let table = super::query_auth_table(); + let err = table.authorize_unrestricted_read().await.unwrap_err(); + assert!( + matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), + "write-path authorization must fail closed, got: {err}" + ); + } + + #[test] + fn test_branch_copy_inherits_query_auth_flag() { + use crate::spec::CoreOptions; + + // `query-auth.enabled` arrives on the REST table response, so a branch's + // on-disk schema need not carry it. If the branch copy dropped it, + // `t$branch_x` would read the same data files raw and unauthorized. + let table = super::query_auth_table(); + assert!(CoreOptions::new(table.schema().options()).query_auth_enabled()); + + // `copy_with_branch` replaces the options wholesale with the branch + // schema's; the flag must survive that replacement. + let branch_schema = + table + .schema() + .copy_with_replaced_options(std::collections::HashMap::from([( + "branch".to_string(), + "b1".to_string(), + )])); + assert!( + !CoreOptions::new(branch_schema.options()).query_auth_enabled(), + "precondition: a replaced-options copy loses the flag on its own" + ); + assert!( + CoreOptions::new(branch_schema.copy_with_query_auth_enabled().options()) + .query_auth_enabled(), + "the branch copy must re-assert query-auth.enabled" + ); + } +} diff --git a/crates/paimon/src/table/partition_stat.rs b/crates/paimon/src/table/partition_stat.rs index 9d78b58a..082cd4ae 100644 --- a/crates/paimon/src/table/partition_stat.rs +++ b/crates/paimon/src/table/partition_stat.rs @@ -64,6 +64,11 @@ impl Table { /// /// Returns an empty Vec when the table has no snapshots yet. pub async fn partition_stats(&self) -> crate::Result> { + // Per-partition record counts and partition-key values are derived from + // the raw manifests, so a row filter would not apply: they disclose the + // size and key space of data the grant hides. Require an unrestricted + // grant (this also covers `list_partitions`, which delegates here). + self.authorize_unrestricted_read().await?; let sm = SnapshotManager::new(self.file_io().clone(), self.location().to_string()); let snapshot = match sm.get_latest_snapshot().await? { Some(s) => s, diff --git a/crates/paimon/src/table/query_auth.rs b/crates/paimon/src/table/query_auth.rs new file mode 100644 index 00000000..45131eb0 --- /dev/null +++ b/crates/paimon/src/table/query_auth.rs @@ -0,0 +1,1085 @@ +// 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. + +//! Query-auth enforcement: apply the REST server's per-user row filter and +//! column masking exactly to read output. Parsing of the Java `Predicate` / +//! `Transform` JSON lives in [`crate::spec`]; anything unrecognised is an +//! error, so callers keep the table fail-closed. + +use crate::arrow::residual::{ + boolean_mask_from_predicate, evaluate_column_predicate, literal_scalar_for_arrow_filter, + sanitize_filter_mask, +}; +use crate::spec::{ + DataField, DataType, Datum, Predicate, PredicateOperator, Transform, TransformInput, +}; +use crate::{Error, Result}; +use arrow_arith::boolean::{and_kleene, not, or_kleene}; +use arrow_array::{ArrayRef, BooleanArray, Float32Array, Float64Array, RecordBatch}; +use std::collections::HashSet; +use std::sync::Arc; + +/// Row filters and column masks the REST server granted the current user for a +/// specific set of columns. `authorized = None` means all columns were approved +/// (the request used `select = all`); `Some(set)` scopes the grant to those +/// table-schema indices. Only the REST catalog constructs grants. +#[derive(Debug, Clone, Default, PartialEq)] +pub(crate) struct QueryAuthGrant { + filters: Vec, + masks: Vec, + authorized: Option>, + /// Schema this grant's positional filter/mask indices were parsed against. + /// Enforcing it on a different schema would bind them to other columns. + schema_id: Option, +} + +impl QueryAuthGrant { + pub(crate) fn new( + filters: Vec, + masks: Vec, + authorized: Option>, + ) -> Self { + Self { + filters, + masks, + authorized, + schema_id: None, + } + } + + /// Bind this grant to the schema its indices were resolved against. + pub(crate) fn with_schema_id(mut self, schema_id: i64) -> Self { + self.schema_id = Some(schema_id); + self + } + + /// Whether this grant may be enforced against `schema_id`. A grant with no + /// binding (test-constructed) matches anything. + pub(crate) fn matches_schema(&self, schema_id: i64) -> bool { + self.schema_id.is_none_or(|bound| bound == schema_id) + } + + /// Fully unrestricted: every column approved, no filter, no masking. + pub(crate) fn is_unrestricted(&self) -> bool { + self.authorized.is_none() && self.filters.is_empty() && self.masks.is_empty() + } + + /// Whether the SERVER restricted this user (a row filter or a column mask), + /// as opposed to the client merely scoping its own projection. Statistics + /// suppression, split transport and the incremental-read gate key off this: + /// a column-scoped grant with no filter and no mask distorts nothing. + pub(crate) fn has_server_restrictions(&self) -> bool { + !self.filters.is_empty() || !self.masks.is_empty() + } + + pub(crate) fn filters(&self) -> &[Predicate] { + &self.filters + } + + pub(crate) fn masks(&self) -> &[ColumnMask] { + &self.masks + } + + /// Whether every table-schema index in `columns` was authorized. A grant + /// scoped to a subset does not authorize columns outside it, so a wider + /// projection or a predicate on an un-approved column fails closed. + pub(crate) fn authorizes_columns(&self, columns: impl IntoIterator) -> bool { + match &self.authorized { + None => true, + Some(set) => columns.into_iter().all(|c| set.contains(&c)), + } + } + + /// Whether this grant carries a row filter. Such a filter is applied as a + /// residual pass in `TableRead::to_arrow`, so split row counts, count + /// statistics, and count-based limit pushdown are no longer exact. + pub(crate) fn has_row_filter(&self) -> bool { + !self.filters.is_empty() + } + + /// Table-schema indices of columns this grant masks (empty when no masking). + /// Callers must not push predicates on these columns to scan pruning, which + /// would leak the raw value via row presence. + pub(crate) fn masked_columns(&self) -> Vec { + self.masks.iter().map(|m| m.column).collect() + } + + /// The first of `columns` (table-schema indices) this grant does not + /// authorize, if any. + pub(crate) fn first_unauthorized( + &self, + columns: impl IntoIterator, + ) -> Option { + columns.into_iter().find(|c| !self.authorizes_columns([*c])) + } + + /// Field IDs the grant must physically read (row-filter columns, mask + /// targets, and mask inputs). Scan projection planning must include these so + /// data-evolution column-slice pruning does not drop files that hold them + /// (an omitted column would read as null and wrongly satisfy `IS_NULL`). + pub(crate) fn read_field_ids(&self, fields: &[DataField]) -> Vec { + let mut indices = HashSet::new(); + for filter in &self.filters { + filter.collect_leaf_field_indices(&mut indices); + } + for mask in &self.masks { + indices.insert(mask.column); + mask.transform.collect_field_indices(&mut indices); + } + indices + .into_iter() + .filter_map(|i| fields.get(i).map(|f| f.id())) + .collect() + } +} + +/// Mask one column (by table-schema index) with a transform. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ColumnMask { + pub(crate) column: usize, + pub(crate) transform: Transform, +} + +fn unsupported(message: String) -> Error { + Error::Unsupported { message } +} + +fn field_name(fields: &[DataField], column: usize) -> &str { + fields.get(column).map(|f| f.name()).unwrap_or("?") +} + +/// Fail-closed error for a caller predicate on a masked column (would leak the +/// raw value via row selection). Shared by every builder/scan choke point. +pub(crate) fn masked_filter_error(fields: &[DataField], column: usize) -> Error { + unsupported(format!( + "cannot filter on masked column `{}`", + field_name(fields, column) + )) +} + +/// Fail-closed error for a read that touches a column outside the grant's scope. +pub(crate) fn unauthorized_column_error(fields: &[DataField], column: usize) -> Error { + unsupported(format!( + "query-auth read touches column `{}` outside the authorized set", + field_name(fields, column) + )) +} + +/// Live query-auth scope check shared by the read/scan gates: fail closed when +/// the caller filter references a masked column (pruning on its raw value would +/// leak it) or when the projection/filter touches a column outside the grant's +/// authorized scope. `projected = None` means all columns. +pub(crate) fn scope_check( + grant: &QueryAuthGrant, + fields: &[DataField], + filter_columns: &HashSet, + projected: Option>, +) -> crate::Result<()> { + if let Some(column) = grant + .masked_columns() + .into_iter() + .find(|c| filter_columns.contains(c)) + { + return Err(masked_filter_error(fields, column)); + } + let projected = projected.unwrap_or_else(|| (0..fields.len()).collect()); + if let Some(column) = + grant.first_unauthorized(projected.into_iter().chain(filter_columns.iter().copied())) + { + return Err(unauthorized_column_error(fields, column)); + } + Ok(()) +} + +/// Parse the auth response's JSON filter strings into predicates whose leaf +/// indices refer to `fields` (table-schema order). Empty strings are skipped +/// (Java parity); anything unparseable is an error. +pub(crate) fn parse_auth_filters( + filters: &[String], + fields: &[DataField], +) -> Result> { + filters + .iter() + .filter(|f| !f.trim().is_empty()) + .map(|f| Predicate::from_rest_json(f, fields)) + .collect() +} + +// ==================== Exact evaluation ==================== + +/// Exactly evaluate the ANDed `predicates` against `batch` (whose columns +/// correspond 1:1 to `batch_fields`; leaf indices refer to `schema_fields`) +/// and drop non-matching rows. Unlike the pruning evaluators, anything that +/// cannot be evaluated is an error — a security filter must not fall open. +pub(crate) fn strict_filter_batch( + batch: &RecordBatch, + predicates: &[Predicate], + schema_fields: &[DataField], + batch_fields: &[DataField], +) -> Result { + let mut combined: Option = None; + for predicate in predicates { + let mask = strict_mask(batch, predicate, schema_fields, batch_fields)?; + combined = Some(match combined { + Some(existing) => kleene(and_kleene(&existing, &mask))?, + None => mask, + }); + } + let Some(mask) = combined else { + return Ok(batch.clone()); + }; + let mask = sanitize_filter_mask(mask); + arrow_select::filter::filter_record_batch(batch, &mask).map_err(|e| Error::DataInvalid { + message: format!("failed to apply query-auth row filter: {e}"), + source: Some(Box::new(e)), + }) +} + +fn strict_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) => fold_masks(batch, children, schema_fields, batch_fields, true), + Predicate::Or(children) => fold_masks(batch, children, schema_fields, batch_fields, false), + Predicate::Not(inner) => { + let mask = strict_mask(batch, inner, schema_fields, batch_fields)?; + kleene(not(&mask)) + } + Predicate::Leaf { + index, + op, + literals, + .. + } => { + let field = schema_fields.get(*index).ok_or_else(|| { + unsupported(format!( + "query-auth filter references unknown field #{index}" + )) + })?; + let position = batch_fields + .iter() + .position(|f| f.id() == field.id() && f.name() == field.name()) + .ok_or_else(|| { + unsupported(format!( + "query-auth filter field `{}` missing from read", + field.name() + )) + })?; + let column = canonicalize_nan(batch.column(position)); + strict_leaf_mask(&column, field.data_type(), *op, literals) + } + } +} + +/// Replace every NaN in a float column with the canonical (positive) NaN. +/// +/// Arrow's comparison kernels use a total ordering in which a NEGATIVE NaN sorts +/// below every finite value, so an auth filter like `f < 0` would admit a +/// negative-NaN row. Java canonicalizes NaN in `Float`/`Double.compare`, making +/// every NaN greater than all finite values, so the row is rejected. Applied +/// only to authorization filters — ordinary query pushdown keeps Arrow +/// semantics. Signed zero already agrees with Java and is left alone. +fn canonicalize_nan(column: &ArrayRef) -> ArrayRef { + match column.data_type() { + arrow_schema::DataType::Float32 => { + let values = column.as_any().downcast_ref::(); + match values { + 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 + } + })) as ArrayRef + } + _ => Arc::clone(column), + } + } + arrow_schema::DataType::Float64 => { + let values = column.as_any().downcast_ref::(); + match values { + 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 + } + })) as ArrayRef + } + _ => Arc::clone(column), + } + } + _ => Arc::clone(column), + } +} + +fn fold_masks( + batch: &RecordBatch, + children: &[Predicate], + schema_fields: &[DataField], + batch_fields: &[DataField], + use_and: bool, +) -> Result { + let mut combined: Option = None; + for child in children { + let mask = strict_mask(batch, child, schema_fields, batch_fields)?; + combined = Some(match combined { + Some(existing) if use_and => kleene(and_kleene(&existing, &mask))?, + Some(existing) => kleene(or_kleene(&existing, &mask))?, + None => mask, + }); + } + combined.ok_or_else(|| unsupported("query-auth filter has an empty compound".to_string())) +} + +fn strict_leaf_mask( + column: &ArrayRef, + data_type: &DataType, + op: PredicateOperator, + literals: &[Datum], +) -> Result { + let scalar = |literal: &Datum| -> Result> { + literal_scalar_for_arrow_filter(literal, data_type)?.ok_or_else(|| { + unsupported(format!( + "query-auth filter literal is not comparable to type {data_type:?}" + )) + }) + }; + match op { + PredicateOperator::IsNull => Ok(boolean_mask_from_predicate(column.len(), |row| { + column.is_null(row) + })), + PredicateOperator::IsNotNull => Ok(boolean_mask_from_predicate(column.len(), |row| { + column.is_valid(row) + })), + PredicateOperator::In | PredicateOperator::NotIn => { + // Kleene IN: OR of equalities; x NOT IN (..) = NOT(IN), nulls stay null. + let mut combined = BooleanArray::from(vec![false; column.len()]); + for literal in literals { + let eq = kleene(evaluate_column_predicate( + column, + &scalar(literal)?, + PredicateOperator::Eq, + ))?; + combined = kleene(or_kleene(&combined, &eq))?; + } + if matches!(op, PredicateOperator::NotIn) { + combined = kleene(not(&combined))?; + if literals.is_empty() { + // `x NOT IN ()` is true only for non-null rows. + combined = + boolean_mask_from_predicate(column.len(), |row| column.is_valid(row)); + } + } + Ok(combined) + } + PredicateOperator::Eq + | PredicateOperator::NotEq + | PredicateOperator::Lt + | PredicateOperator::LtEq + | PredicateOperator::Gt + | PredicateOperator::GtEq + | PredicateOperator::StartsWith + | PredicateOperator::EndsWith + | PredicateOperator::Contains + | PredicateOperator::Like => { + let literal = literals.first().ok_or_else(|| { + unsupported("query-auth filter comparison without literal".to_string()) + })?; + kleene(evaluate_column_predicate(column, &scalar(literal)?, op)) + } + PredicateOperator::Between | PredicateOperator::NotBetween => { + let (Some(low), Some(high)) = (literals.first(), literals.get(1)) else { + return Err(unsupported( + "query-auth BETWEEN filter without bounds".to_string(), + )); + }; + let lo = kleene(evaluate_column_predicate( + column, + &scalar(low)?, + PredicateOperator::GtEq, + ))?; + let hi = kleene(evaluate_column_predicate( + column, + &scalar(high)?, + PredicateOperator::LtEq, + ))?; + let between = kleene(and_kleene(&lo, &hi))?; + if matches!(op, PredicateOperator::NotBetween) { + kleene(not(&between)) + } else { + Ok(between) + } + } + } +} + +fn kleene( + result: std::result::Result, +) -> Result { + result.map_err(|e| Error::DataInvalid { + message: format!("failed to evaluate query-auth row filter: {e}"), + source: Some(Box::new(e)), + }) +} + +// ==================== Column masking ==================== + +fn mask_err(detail: impl std::fmt::Display) -> Error { + unsupported(format!("cannot parse query-auth column masking: {detail}")) +} + +/// Parse the auth response's `columnMasking` map (column name -> Java +/// `Transform` JSON) against `fields` (table-schema order). +pub(crate) fn parse_column_masking( + masking: &std::collections::HashMap, + fields: &[DataField], +) -> Result> { + let mut masks = Vec::with_capacity(masking.len()); + for (column, json) in masking { + let target = fields + .iter() + .position(|f| f.name() == column) + .ok_or_else(|| mask_err(format!("unknown field `{column}`")))?; + let transform = Transform::from_rest_json(json, fields)?; + // The masked value replaces the column in place, so its type must match + // the column's; a type-changing transform (e.g. `CAST(id AS STRING)` on + // an INT column) cannot be represented and must fail closed rather than + // be cast back to the raw type. Types are compared via their arrow + // representation (nullability-agnostic). + if let Some(out) = mask_output_type(&transform, fields) { + let target_type = crate::arrow::paimon_type_to_arrow(fields[target].data_type())?; + if out != target_type { + return Err(mask_err(format!( + "masking `{column}` produces {out:?} but the column is {target_type:?}" + ))); + } + } + // A mask that can yield null on a NOT NULL column would leave the output + // batch's schema claiming non-nullable, letting engines fold + // `col IS [NOT] NULL` before the masked-predicate guard. Fail closed. + if !fields[target].data_type().is_nullable() && mask_can_be_null(&transform, fields) { + return Err(mask_err(format!( + "masking `{column}` can produce null but the column is NOT NULL" + ))); + } + masks.push(ColumnMask { + column: target, + transform, + }); + } + // Deterministic order regardless of map iteration. + masks.sort_by_key(|m| m.column); + + // Masks read their inputs from the RAW batch (like Java), so one mask that + // references ANOTHER mask's target would copy that column's unmasked value + // into its own output — defeating the second mask. Referencing your own + // target is the normal case (`name := UPPER(name)`) and stays valid. + let targets: HashSet = masks.iter().map(|m| m.column).collect(); + for mask in &masks { + let mut inputs = HashSet::new(); + mask.transform.collect_field_indices(&mut inputs); + if let Some(other) = inputs + .into_iter() + .find(|i| *i != mask.column && targets.contains(i)) + { + return Err(mask_err(format!( + "masking `{}` reads masked column `{}`, which would expose its raw value", + field_name(fields, mask.column), + field_name(fields, other) + ))); + } + } + Ok(masks) +} + +/// Whether a mask transform can produce a null value. +fn mask_can_be_null(transform: &Transform, fields: &[DataField]) -> bool { + let field_nullable = |index: &usize| fields[*index].data_type().is_nullable(); + let input_nullable = |inputs: &[TransformInput]| { + inputs.iter().any(|i| match i { + TransformInput::Literal(literal) => literal.is_none(), + TransformInput::Field(index) => field_nullable(index), + }) + }; + match transform { + Transform::Null => true, + Transform::FieldRef(index) | Transform::Cast(index, _) => field_nullable(index), + Transform::Upper(inputs) | Transform::Lower(inputs) | Transform::Concat(inputs) => { + input_nullable(inputs) + } + // CONCAT_WS skips null payloads, so only a null separator (the first + // input) can make the result null. + Transform::ConcatWs(inputs) => inputs.first().is_some_and(|sep| match sep { + TransformInput::Literal(literal) => literal.is_none(), + TransformInput::Field(index) => field_nullable(index), + }), + } +} + +/// Arrow output type of a mask transform, or `None` when it always matches the +/// target column (the `NULL` transform builds a null of the column's own type). +fn mask_output_type(transform: &Transform, fields: &[DataField]) -> Option { + let of = |index: &usize| crate::arrow::paimon_type_to_arrow(fields[*index].data_type()).ok(); + match transform { + Transform::Null => None, + Transform::FieldRef(index) => of(index), + Transform::Cast(_, to) => crate::arrow::paimon_type_to_arrow(to).ok(), + Transform::Upper(_) + | Transform::Lower(_) + | Transform::Concat(_) + | Transform::ConcatWs(_) => Some(arrow_schema::DataType::Utf8), + } +} + +/// Overwrite masked columns of `batch` (whose columns correspond 1:1 to +/// `batch_fields`). Masks whose target column is not in the batch are skipped +/// (Java parity); anything that cannot be evaluated is an error. +pub(crate) fn mask_batch( + batch: &RecordBatch, + masks: &[ColumnMask], + schema_fields: &[DataField], + batch_fields: &[DataField], +) -> Result { + use arrow_array::new_null_array; + + // Nothing to mask (e.g. a `COUNT(*)` read projects no columns); return the + // batch unchanged, preserving its row count even with zero columns. + if masks.is_empty() { + return Ok(batch.clone()); + } + + // All batch positions holding a given table-schema field (a projection may + // repeat a column, so every copy must be masked, not just the first). + let positions_of = |schema_index: usize| -> Result> { + let field = schema_fields.get(schema_index).ok_or_else(|| { + unsupported(format!( + "query-auth mask references unknown field #{schema_index}" + )) + })?; + // Match on field id alone: the read type may carry the column under a + // different name (projection aliasing, or a renamed column read from an + // older file schema). Also matching the name would silently find no + // position and skip the mask, emitting the raw value. + Ok(batch_fields + .iter() + .enumerate() + .filter(|(_, f)| f.id() == field.id()) + .map(|(pos, _)| pos) + .collect()) + }; + let input_column = |schema_index: usize| -> Result { + positions_of(schema_index)? + .first() + .map(|pos| batch.column(*pos).clone()) + .ok_or_else(|| unsupported("query-auth mask input missing from read".to_string())) + }; + + let mut columns = batch.columns().to_vec(); + for mask in masks { + let targets = positions_of(mask.column)?; + // `masks` is already filtered to targets the caller projects, so a mask + // that resolves to no batch column means the read cannot enforce it — + // emitting the column unmasked would leak the raw value. + let Some(&first) = targets.first() else { + return Err(unsupported(format!( + "query-auth mask for column `{}` cannot be applied: the column is missing \ + from the read", + field_name(schema_fields, mask.column) + ))); + }; + let target_type = batch.schema().field(first).data_type().clone(); + let masked: ArrayRef = match &mask.transform { + Transform::Null => new_null_array(&target_type, batch.num_rows()), + Transform::FieldRef(index) => input_column(*index)?, + Transform::Cast(index, to) => { + let to_arrow = crate::arrow::paimon_type_to_arrow(to)?; + cast_masked(&input_column(*index)?, &to_arrow)? + } + Transform::Upper(inputs) => string_mask(batch, inputs, &input_column, |v| { + Some(v.first()?.as_ref().map(|s| s.to_uppercase())) + })?, + Transform::Lower(inputs) => string_mask(batch, inputs, &input_column, |v| { + Some(v.first()?.as_ref().map(|s| s.to_lowercase())) + })?, + // SQL semantics: CONCAT is null if any input is null; CONCAT_WS + // uses the first input as separator and skips null values. + Transform::Concat(inputs) => string_mask(batch, inputs, &input_column, |v| { + Some( + v.iter() + .cloned() + .collect::>>() + .map(|p| p.concat()), + ) + })?, + Transform::ConcatWs(inputs) => string_mask(batch, inputs, &input_column, |v| { + let (sep, rest) = v.split_first()?; + Some( + sep.as_ref() + .map(|sep| rest.iter().flatten().cloned().collect::>().join(sep)), + ) + })?, + }; + // Parse-time type checks guarantee a compatible type, so this only + // aligns arrow representations. Mask every copy of the target column. + let masked = cast_masked(&masked, &target_type)?; + for pos in targets { + columns[pos] = masked.clone(); + } + } + RecordBatch::try_new_with_options( + batch.schema(), + columns, + &arrow_array::RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| Error::DataInvalid { + message: format!("failed to apply query-auth column masking: {e}"), + source: Some(Box::new(e)), + }) +} + +fn cast_masked(array: &ArrayRef, to: &arrow_schema::DataType) -> Result { + if array.data_type() == to { + return Ok(array.clone()); + } + arrow_cast::cast(array, to).map_err(|e| { + unsupported(format!( + "query-auth mask value of type {:?} cannot be cast to column type {to:?}: {e}", + array.data_type() + )) + }) +} + +/// Evaluate a string transform row by row. `combine` receives the resolved +/// inputs (None = SQL NULL) and returns the masked value; a `None` from +/// `combine` means the transform is malformed for this input arity. +fn string_mask( + batch: &RecordBatch, + inputs: &[TransformInput], + input_column: &dyn Fn(usize) -> Result, + combine: impl Fn(&[Option]) -> Option>, +) -> Result { + use arrow_array::{Array, StringArray}; + use std::sync::Arc; + + // Resolve field inputs once, as string arrays (None = literal slot). + let resolved: Vec> = inputs + .iter() + .map(|input| match input { + TransformInput::Literal(_) => Ok(None), + TransformInput::Field(index) => { + cast_masked(&input_column(*index)?, &arrow_schema::DataType::Utf8).map(Some) + } + }) + .collect::>>()?; + + let mut values: Vec> = Vec::with_capacity(batch.num_rows()); + for row in 0..batch.num_rows() { + let row_inputs = inputs + .iter() + .zip(&resolved) + .map(|(input, array)| match (input, array) { + (TransformInput::Literal(s), _) => Ok(s.clone()), + (TransformInput::Field(_), Some(array)) => { + let strings = array + .as_any() + .downcast_ref::() + .ok_or_else(|| unsupported("mask input is not a string".to_string()))?; + Ok((!strings.is_null(row)).then(|| strings.value(row).to_string())) + } + (TransformInput::Field(_), None) => unreachable!("field inputs are resolved above"), + }) + .collect::>>()?; + let value = combine(&row_inputs) + .ok_or_else(|| unsupported("query-auth string mask is malformed".to_string()))?; + values.push(value); + } + Ok(Arc::new(StringArray::from(values)) as ArrayRef) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{IntType, VarCharType}; + use arrow_array::{Int32Array, StringArray}; + use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; + use std::sync::Arc; + + fn fields() -> Vec { + vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "name".to_string(), + DataType::VarChar(VarCharType::new(255).unwrap()), + ), + ] + } + + fn leaf_json(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 batch() -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])), + Arc::new(StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + ])), + ], + ) + .unwrap() + } + + #[test] + fn test_strict_filter_batch_filters_rows() { + let fields = fields(); + let filters = + parse_auth_filters(&[leaf_json("GREATER_THAN", "id", "[1]")], &fields).unwrap(); + let filtered = strict_filter_batch(&batch(), &filters, &fields, &fields).unwrap(); + // id > 1 keeps rows 2 and 4; the NULL row is excluded. + assert_eq!(filtered.num_rows(), 2); + } + + #[test] + fn test_strict_filter_matches_java_nan_ordering() { + use crate::spec::{DoubleType, PredicateBuilder}; + use arrow_array::Float64Array; + use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; + + // Java `Double.compare` canonicalizes NaN above every finite value, so + // `f < 0` must reject NaN — including a NEGATIVE NaN, which Arrow's + // total ordering would otherwise sort below finite values and admit. + let fields = vec![DataField::new( + 0, + "f".to_string(), + DataType::Double(DoubleType::new()), + )]; + let neg_nan = f64::from_bits(0xFFF8_0000_0000_0000); + let batch = RecordBatch::try_new( + std::sync::Arc::new(ArrowSchema::new(vec![ArrowField::new( + "f", + ArrowDataType::Float64, + true, + )])), + vec![std::sync::Arc::new(Float64Array::from(vec![ + neg_nan, + f64::NAN, + -1.0, + 1.0, + ]))], + ) + .unwrap(); + + let less = PredicateBuilder::new(&fields) + .less_than("f", crate::spec::Datum::Double(0.0)) + .unwrap(); + let filtered = strict_filter_batch(&batch, &[less], &fields, &fields).unwrap(); + assert_eq!( + filtered.num_rows(), + 1, + "only -1.0 may pass `f < 0`; neither NaN sign may" + ); + + let greater = PredicateBuilder::new(&fields) + .greater_than("f", crate::spec::Datum::Double(0.0)) + .unwrap(); + let filtered = strict_filter_batch(&batch, &[greater], &fields, &fields).unwrap(); + assert_eq!( + filtered.num_rows(), + 3, + "both NaNs and 1.0 pass `f > 0` (NaN is greatest, like Java)" + ); + } + + #[test] + fn test_strict_filter_not_excludes_nulls() { + let fields = fields(); + // NOT (id = 2): NULL rows must stay excluded (SQL three-valued logic). + let json = format!( + r#"{{"kind":"COMPOUND","function":"AND","children":[{}]}}"#, + leaf_json("NOT_EQUAL", "id", "[2]") + ); + let filters = parse_auth_filters(&[json], &fields).unwrap(); + let filtered = strict_filter_batch(&batch(), &filters, &fields, &fields).unwrap(); + assert_eq!( + filtered.num_rows(), + 2, + "rows 1 and 4 only, not the NULL row" + ); + } + + /// Java #7034 baseline: filter each supported literal type exactly. + #[test] + fn test_strict_filter_batch_typed_matrix() { + use crate::spec::{BigIntType, BooleanType, DoubleType, FloatType}; + use arrow_array::{BooleanArray, Float32Array, Float64Array, Int64Array}; + + let fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "age".to_string(), DataType::BigInt(BigIntType::new())), + DataField::new(2, "salary".to_string(), DataType::Double(DoubleType::new())), + DataField::new( + 3, + "is_active".to_string(), + DataType::Boolean(BooleanType::new()), + ), + DataField::new(4, "score".to_string(), DataType::Float(FloatType::new())), + DataField::new( + 5, + "name".to_string(), + DataType::VarChar(VarCharType::new(255).unwrap()), + ), + ]; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("age", ArrowDataType::Int64, true), + ArrowField::new("salary", ArrowDataType::Float64, true), + ArrowField::new("is_active", ArrowDataType::Boolean, true), + ArrowField::new("score", ArrowDataType::Float32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(Int64Array::from(vec![25, 30, 35, 28])), + Arc::new(Float64Array::from(vec![50000.0, 60000.0, 70000.0, 55000.0])), + Arc::new(BooleanArray::from(vec![true, false, true, true])), + Arc::new(Float32Array::from(vec![85.5, 90.0, 95.5, 88.0])), + Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie", "David"])), + ], + ) + .unwrap(); + fn typed_leaf(function: &str, field: &str, ftype: &str, literals: &str) -> String { + format!( + r#"{{"kind":"LEAF","transform":{{"name":"FIELD_REF","fieldRef":{{"index":0,"name":"{field}","type":"{ftype}"}}}},"function":"{function}","literals":{literals}}}"# + ) + } + // (filter, expected surviving ids) — mirrors Java MockRESTCatalogTest. + let cases: Vec<(String, Vec)> = vec![ + (typed_leaf("GREATER_THAN", "id", "INT", "[2]"), vec![3, 4]), + ( + typed_leaf("GREATER_OR_EQUAL", "age", "BIGINT", "[30]"), + vec![2, 3], + ), + ( + typed_leaf("GREATER_THAN", "salary", "DOUBLE", "[55000.0]"), + vec![2, 3], + ), + ( + typed_leaf("EQUAL", "is_active", "BOOLEAN", "[true]"), + vec![1, 3, 4], + ), + ( + typed_leaf("GREATER_OR_EQUAL", "score", "FLOAT", "[90.0]"), + vec![2, 3], + ), + ( + typed_leaf("EQUAL", "name", "STRING", "[\"Alice\"]"), + vec![1], + ), + ( + // Two predicates ANDed by the grant list semantics. + format!( + r#"{{"kind":"COMPOUND","function":"AND","children":[{},{}]}}"#, + typed_leaf("GREATER_OR_EQUAL", "age", "BIGINT", "[30]"), + typed_leaf("EQUAL", "is_active", "BOOLEAN", "[true]") + ), + vec![3], + ), + ]; + for (json, expected) in cases { + let filters = parse_auth_filters(std::slice::from_ref(&json), &fields).unwrap(); + let filtered = strict_filter_batch(&batch, &filters, &fields, &fields).unwrap(); + let ids = filtered + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ids, expected, "for {json}"); + } + } + + fn masking(column: &str, json: &str) -> std::collections::HashMap { + std::collections::HashMap::from([(column.to_string(), json.to_string())]) + } + + #[test] + fn test_reject_mask_reading_another_masked_column() { + // `name := UPPER(name)` is the normal self-reference and stays valid. + let fields = fields(); + assert!(parse_column_masking( + &masking( + "name", + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + ), + &fields + ) + .is_ok()); + + // But a mask that reads ANOTHER masked column would copy that column's + // raw value out (masks read the unmasked batch), defeating its mask. + let both = std::collections::HashMap::from([ + ("id".to_string(), r#"{"name":"NULL"}"#.to_string()), + ( + "name".to_string(), + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + .to_string(), + ), + ]); + assert!( + parse_column_masking(&both, &fields).is_ok(), + "unrelated masks are fine" + ); + + let cross = std::collections::HashMap::from([ + ("name".to_string(), r#"{"name":"NULL"}"#.to_string()), + ( + "alias".to_string(), + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + .to_string(), + ), + ]); + let mut fields_with_alias = fields.clone(); + fields_with_alias.push(DataField::new( + 2, + "alias".to_string(), + DataType::VarChar(VarCharType::new(255).unwrap()), + )); + let Err(err) = parse_column_masking(&cross, &fields_with_alias) else { + panic!("a mask reading another masked column must fail closed"); + }; + assert!(err.to_string().contains("raw value"), "got: {err}"); + } + + #[test] + fn test_parse_column_masking() { + let fields = fields(); + let masks = parse_column_masking(&masking("name", r#"{"name":"NULL"}"#), &fields).unwrap(); + assert!(matches!(masks[0].transform, Transform::Null)); + assert_eq!(masks[0].column, 1); + + let upper = r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"#; + let masks = parse_column_masking(&masking("name", upper), &fields).unwrap(); + assert!(matches!(&masks[0].transform, Transform::Upper(inputs) if inputs.len() == 1)); + + // Unknown transform / unknown column / bad JSON: all fail closed. + for (column, json) in [ + ("name", r#"{"name":"ROT13"}"#), + ("missing", r#"{"name":"NULL"}"#), + ("name", "not json"), + ] { + assert!(parse_column_masking(&masking(column, json), &fields).is_err()); + } + } + + #[test] + fn test_mask_batch_null_and_string_transforms() { + use arrow_array::Array; + let fields = fields(); + + // NULL mask: the whole column becomes null. + let masks = parse_column_masking(&masking("name", r#"{"name":"NULL"}"#), &fields).unwrap(); + let masked = mask_batch(&batch(), &masks, &fields, &fields).unwrap(); + assert_eq!(masked.column(1).null_count(), 4); + assert_eq!(masked.column(0).null_count(), 1, "other columns untouched"); + + // CONCAT_WS("-", literal, field): "x-a", "x-b", ... + let concat = + r#"{"name":"CONCAT_WS","inputs":["-","x",{"index":1,"name":"name","type":"STRING"}]}"#; + let masks = parse_column_masking(&masking("name", concat), &fields).unwrap(); + let masked = mask_batch(&batch(), &masks, &fields, &fields).unwrap(); + let names = masked + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(names.value(0), "x-a"); + + // Masked column absent from the batch: fail closed. The caller filters + // masks to projected targets, so this means the read cannot enforce the + // mask — emitting the column unmasked would leak the raw value. + let one_col = batch().project(&[0]).unwrap(); + let one_field = vec![fields[0].clone()]; + let masks = parse_column_masking(&masking("name", r#"{"name":"NULL"}"#), &fields).unwrap(); + let Err(err) = mask_batch(&one_col, &masks, &fields, &one_field) else { + panic!("a mask whose target is missing from the read must fail closed"); + }; + assert!(err.to_string().contains("cannot be applied"), "got: {err}"); + } + + #[test] + fn test_reject_type_changing_cast_mask() { + let fields = fields(); + // CAST(id AS STRING) on an INT column changes type -> fail closed. + let cast = + r#"{"name":"CAST","fieldRef":{"index":0,"name":"id","type":"INT"},"type":"STRING"}"#; + assert!(parse_column_masking(&masking("id", cast), &fields).is_err()); + // A string transform on a non-string column also fails closed. + let upper = r#"{"name":"UPPER","inputs":[{"index":0,"name":"id","type":"INT"}]}"#; + assert!(parse_column_masking(&masking("id", upper), &fields).is_err()); + } + + #[test] + fn test_mask_batch_masks_every_duplicate_target() { + use arrow_array::Array; + let fields = fields(); + let masks = parse_column_masking(&masking("name", r#"{"name":"NULL"}"#), &fields).unwrap(); + // A projection that repeats the masked column: both copies must be masked. + let base = batch(); + let dup = base.project(&[1, 1]).unwrap(); + let dup_fields = vec![fields[1].clone(), fields[1].clone()]; + let masked = mask_batch(&dup, &masks, &fields, &dup_fields).unwrap(); + assert_eq!(masked.column(0).null_count(), 4); + assert_eq!(masked.column(1).null_count(), 4); + } + + #[test] + fn test_grant_authorized_column_scope() { + // A grant scoped to a subset is not globally unrestricted and rejects + // columns outside the approved set. + let grant = QueryAuthGrant::new(Vec::new(), Vec::new(), Some(HashSet::from([0]))); + assert!(!grant.is_unrestricted()); + assert!(grant.authorizes_columns([0])); + assert!(!grant.authorizes_columns([1])); + // `None` (all columns) with no filter/mask is fully unrestricted. + let all = QueryAuthGrant::new(Vec::new(), Vec::new(), None); + assert!(all.is_unrestricted()); + assert!(all.authorizes_columns([0, 1, 99])); + } +} diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index de432a5c..5d3a3d60 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -291,6 +291,10 @@ struct PaimonReadBuilder<'a> { limit: Option, row_ranges: Option>, case_sensitive: bool, + /// Table-schema indices referenced by the full caller filter (before it is + /// split into partition/data conjuncts). The query-auth gates check these + /// against the grant fetched at plan time. + filter_columns: HashSet, } impl<'a> PaimonReadBuilder<'a> { @@ -303,6 +307,7 @@ impl<'a> PaimonReadBuilder<'a> { limit: None, row_ranges: None, case_sensitive: true, + filter_columns: HashSet::new(), } } @@ -362,6 +367,12 @@ impl<'a> PaimonReadBuilder<'a> { /// primary-key merge reads push key conjuncts below the merge and enforce /// the full predicate with an exact post-merge residual filter. pub fn with_filter(&mut self, filter: Predicate) -> &mut Self { + // Capture the columns of the FULL predicate before it is split into + // partition/data conjuncts, so both the masked-column guard and the + // authorized-scope check see a masked/unauthorized partition key (which + // would otherwise be pruned on its raw value). + self.filter_columns.clear(); + filter.collect_leaf_field_indices(&mut self.filter_columns); self.filter = normalize_filter(self.table, filter); self.try_extract_row_id_ranges(); self @@ -434,6 +445,10 @@ impl<'a> PaimonReadBuilder<'a> { let partition_filter = self.filter.partition_predicate.clone().map(|pred| { PartitionFilter::from_predicate(pred, &self.table.schema().partition_fields()) }); + // The grant's auth field IDs are folded into the scan projection inside + // `TableScan::plan` — where the grant has been fetched — not here, where + // it does not yet exist (that early read of an empty grant was a + // fail-open row-filter bypass). let read_type = self.resolve_read_type().unwrap_or(None); TableScan::new( self.table, @@ -448,14 +463,35 @@ impl<'a> PaimonReadBuilder<'a> { &self.filter.data_predicates, self.table.schema().fields(), )) + .with_query_auth_scope(self.filter_columns.clone(), self.projected_schema_indices()) + } + + /// Table-schema indices of the projected columns (`None` = all). + fn projected_schema_indices(&self) -> Option> { + // Resolve names too: a `with_projection` selection lives in + // `projection_names`, and scoping the grant to all columns would deny a + // user authorized for exactly the requested subset. An unresolvable + // projection falls back to the full scope; `new_read` reports the error. + self.resolve_read_type().ok().flatten().map(|fields| { + fields + .iter() + .filter_map(|f| { + self.table + .schema() + .fields() + .iter() + .position(|s| s.id() == f.id()) + }) + .collect() + }) } /// Create a table read for consuming splits (e.g. from a scan plan). pub fn new_read(&self) -> Result> { - // Fail closed at read construction so bindings that short-circuit before - // `to_arrow` (e.g. an empty-splits fast path) can't bypass the guard. - let core_options = self.table.schema.core_options(); - core_options.ensure_read_authorized()?; + // Query-auth is enforced in `TableRead::to_arrow` off the grant stamped + // on the splits by planning (fail closed when a query-auth table's + // splits carry no grant), so no gate is needed at read construction — + // an empty-splits fast path produces no rows to leak. let read_type = match self.resolve_read_type()? { None => self.table.schema.fields().to_vec(), Some(fields) => fields, @@ -813,31 +849,237 @@ mod tests { .unwrap() } - #[test] - fn test_read_fails_closed_when_query_auth_enabled() { + /// A real split carrying no query-auth grant (an unauthorized read path). + fn ungranted_split() -> crate::table::DataSplit { + DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("file:/tmp/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(vec![test_data_file("data.parquet", 4, 1)]) + .build() + .unwrap() + } + + #[tokio::test] + async fn test_read_fails_closed_when_query_auth_enabled() { let table = query_auth_table(); - // `new_read` fails closed, so bindings that short-circuit before `to_arrow` can't bypass. - let err = table.new_read_builder().new_read().unwrap_err(); + // Enforcement is at `to_arrow` off the split grant: a read whose splits + // carry no grant (never authorized by planning) must fail closed, so + // bindings that short-circuit can't bypass. + let read = table.new_read_builder().new_read().unwrap(); + let Err(err) = read.to_arrow(&[ungranted_split()]) else { + panic!("a query-auth read without a stamped grant must fail closed"); + }; assert!( matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), "building a read for a query-auth.enabled table must fail closed" ); } - #[test] - fn test_dynamic_option_cannot_disable_query_auth() { + #[tokio::test] + async fn test_dynamic_option_cannot_disable_query_auth() { // Copying the table with the option off must not weaken a stored `true`. let table = query_auth_table().copy_with_options(HashMap::from([( "query-auth.enabled".to_string(), "false".to_string(), )])); - let err = table.new_read_builder().new_read().unwrap_err(); + let read = table.new_read_builder().new_read().unwrap(); + let Err(err) = read.to_arrow(&[ungranted_split()]) else { + panic!("a dynamic override must not disable query-auth"); + }; assert!( matches!(err, crate::Error::Unsupported { ref message } if message.contains("query-auth.enabled")), "a dynamic override must not disable query-auth" ); } + #[tokio::test] + async fn test_query_auth_filtered_grant_filters_rows_exactly() { + let tempdir = tempdir().unwrap(); + let table_path = local_file_path(tempdir.path()); + let bucket_dir = tempdir.path().join("bucket-0"); + fs::create_dir_all(&bucket_dir).unwrap(); + + let parquet_path = bucket_dir.join("data.parquet"); + write_int_parquet_file( + &parquet_path, + vec![("id", vec![1, 2, 3, 4]), ("value", vec![1, 2, 20, 30])], + None, + ); + let file_size = fs::metadata(&parquet_path).unwrap().len() as i64; + + let file_io = FileIOBuilder::new("file").build().unwrap(); + let table_schema = TableSchema::new( + 0, + &Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .option("query-auth.enabled", "true") + .build() + .unwrap(), + ); + let table = Table::new( + file_io, + Identifier::new("default", "t"), + table_path, + table_schema, + None, + ); + // Grant: the user may only see rows with value >= 10. The filter column + // is NOT in the projection, so the read must fetch it and project it away. + // The grant is threaded on the split (as scan planning would stamp it). + let auth_filter = PredicateBuilder::new(table.schema().fields()) + .greater_or_equal("value", crate::spec::Datum::Int(10)) + .unwrap(); + let grant = std::sync::Arc::new(crate::table::query_auth::QueryAuthGrant::new( + vec![auth_filter], + Vec::new(), + None, + )); + + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(local_file_path(&bucket_dir)) + .with_total_buckets(1) + .with_data_files(vec![test_data_file("data.parquet", 4, file_size)]) + .build() + .unwrap() + .with_query_auth_grant(Some(grant)); + + let read = TableRead::new(&table, vec![table.schema().fields()[0].clone()], Vec::new()); + let batches = read + .to_arrow(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(collect_int_column(&batches, "id"), vec![3, 4]); + // The filter column must not leak into the output schema. + assert_eq!(batches[0].num_columns(), 1); + } + + #[tokio::test] + async fn test_query_auth_masked_grant_masks_and_guards_predicates() { + use crate::table::query_auth::{parse_column_masking, QueryAuthGrant}; + use arrow_array::Array; + + let tempdir = tempdir().unwrap(); + let table_path = local_file_path(tempdir.path()); + let bucket_dir = tempdir.path().join("bucket-0"); + fs::create_dir_all(&bucket_dir).unwrap(); + let parquet_path = bucket_dir.join("data.parquet"); + write_int_parquet_file( + &parquet_path, + vec![("id", vec![1, 2, 3, 4]), ("value", vec![1, 2, 20, 30])], + None, + ); + let file_size = fs::metadata(&parquet_path).unwrap().len() as i64; + + let file_io = FileIOBuilder::new("file").build().unwrap(); + let table_schema = TableSchema::new( + 0, + &Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("value", DataType::Int(IntType::new())) + .option("query-auth.enabled", "true") + .build() + .unwrap(), + ); + let table = Table::new( + file_io, + Identifier::new("default", "t"), + table_path, + table_schema, + None, + ); + // Grant: filter on raw `value` >= 10, then mask `value` with NULL. + let auth_filter = PredicateBuilder::new(table.schema().fields()) + .greater_or_equal("value", crate::spec::Datum::Int(10)) + .unwrap(); + let masks = parse_column_masking( + &std::collections::HashMap::from([( + "value".to_string(), + r#"{"name":"NULL"}"#.to_string(), + )]), + table.schema().fields(), + ) + .unwrap(); + let grant = std::sync::Arc::new(QueryAuthGrant::new(vec![auth_filter], masks, None)); + + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(local_file_path(&bucket_dir)) + .with_total_buckets(1) + .with_data_files(vec![test_data_file("data.parquet", 4, file_size)]) + .build() + .unwrap() + .with_query_auth_grant(Some(grant)); + + // Filter runs on raw values, then the surviving rows are masked. + let read = TableRead::new(&table, table.schema().fields().to_vec(), Vec::new()); + let batches = read + .to_arrow(std::slice::from_ref(&split)) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_int_column(&batches, "id"), vec![3, 4]); + assert_eq!(batches[0].column(1).null_count(), 2, "value masked to NULL"); + + // A caller predicate on the masked column must fail closed (oracle guard). + let caller_filter = PredicateBuilder::new(table.schema().fields()) + .equal("value", crate::spec::Datum::Int(20)) + .unwrap(); + let read = TableRead::new( + &table, + table.schema().fields().to_vec(), + vec![caller_filter], + ); + let Err(err) = read.to_arrow(&[split]) else { + panic!("filtering on a masked column must fail closed"); + }; + assert!(err.to_string().contains("masked column"), "got: {err}"); + } + + #[tokio::test] + async fn test_query_auth_scope_rejects_unauthorized_column() { + use crate::table::query_auth::QueryAuthGrant; + // A grant scoped to no columns must fail closed when the read projects + // `id` — the scope check runs in `to_arrow` before any data is read + // (and also at plan time; see the rest_catalog integration test). + let table = query_auth_table(); + let grant = std::sync::Arc::new(QueryAuthGrant::new( + Vec::new(), + Vec::new(), + Some(HashSet::new()), + )); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("/tmp/does-not-matter".to_string()) + .with_total_buckets(1) + .with_data_files(vec![test_data_file("data.parquet", 4, 1)]) + .build() + .unwrap() + .with_query_auth_grant(Some(grant)); + let read = TableRead::new(&table, vec![table.schema().fields()[0].clone()], Vec::new()); + let Err(err) = read.to_arrow(&[split]) else { + panic!("reading an unauthorized column must fail closed"); + }; + assert!( + err.to_string().contains("outside the authorized set"), + "got: {err}" + ); + } + #[test] fn test_projected_read_field_ids_uses_projection_ids() { let read_type = vec![DataField::new( diff --git a/crates/paimon/src/table/rest_env.rs b/crates/paimon/src/table/rest_env.rs index 4f2a9df7..1e29a632 100644 --- a/crates/paimon/src/table/rest_env.rs +++ b/crates/paimon/src/table/rest_env.rs @@ -188,6 +188,15 @@ impl RESTEnv { )) } + /// Fetch the per-user row filter and column masking for this table. + /// Mirrors Java `CatalogEnvironment.tableQueryAuth()`. + pub(crate) async fn table_query_auth( + &self, + select: Option>, + ) -> Result { + self.api.auth_table_query(&self.identifier, select).await + } + /// Create a `RESTSnapshotCommit` from this environment. pub fn snapshot_commit(&self) -> Arc { Arc::new(RESTSnapshotCommit::new( diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 41836580..9f296a2d 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -20,8 +20,10 @@ //! Reference: [org.apache.paimon.table.source](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/table/source/). use crate::spec::{BinaryRow, DataFileMeta}; +use crate::table::query_auth::QueryAuthGrant; use crate::table::stats_filter::group_by_overlapping_row_id; use serde::{Deserialize, Serialize}; +use std::sync::Arc; fn is_vector_store_file_name(file_name: &str) -> bool { file_name.to_ascii_lowercase().contains(".vector.") @@ -471,7 +473,7 @@ impl PartitionBucket { /// Input split for reading: partition + bucket + list of data files and optional deletion files. /// /// Reference: [org.apache.paimon.table.source.DataSplit](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java) -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Deserialize)] pub struct DataSplit { snapshot_id: i64, partition: BinaryRow, @@ -487,6 +489,54 @@ pub struct DataSplit { /// physical rows are exactly its logical rows (modulo deletion files). /// Mirrors Java `DataSplit#rawConvertible`. raw_convertible: bool, + /// The per-user query-auth grant this split must be read under, stamped by + /// scan planning (or a write-path authorizer). Runtime-only — never + /// serialized — mirroring Java `QueryAuthSplit(split, authResult)`. Carried + /// on the split so `TableRead::to_arrow` enforces the exact grant the plan + /// fetched, instead of a shared mutable slot on `Table`. + #[serde(skip)] + query_auth_grant: Option>, +} + +/// Hand-written so that serializing a split planned under a restricted +/// query-auth grant FAILS instead of silently dropping the grant: no serde +/// format carries it, so the receiver would read the split's files raw. The +/// emitted shape is identical to the derived one (the grant is never a field). +impl Serialize for DataSplit { + fn serialize( + &self, + serializer: S, + ) -> std::result::Result { + use serde::ser::Error as _; + self.ensure_no_restricted_grant("serialize") + .map_err(S::Error::custom)?; + + #[derive(Serialize)] + struct Wire<'a> { + snapshot_id: &'a i64, + partition: &'a BinaryRow, + bucket: &'a i32, + bucket_path: &'a String, + total_buckets: &'a i32, + data_files: &'a Vec, + data_deletion_files: &'a Option>>, + row_ranges: &'a Option>, + raw_convertible: &'a bool, + } + + Wire { + snapshot_id: &self.snapshot_id, + partition: &self.partition, + bucket: &self.bucket, + bucket_path: &self.bucket_path, + total_buckets: &self.total_buckets, + data_files: &self.data_files, + data_deletion_files: &self.data_deletion_files, + row_ranges: &self.row_ranges, + raw_convertible: &self.raw_convertible, + } + .serialize(serializer) + } } impl DataSplit { @@ -539,6 +589,28 @@ impl DataSplit { .all(|file| file.level != 0 && file.delete_row_count == Some(0)) } + /// The query-auth grant this split is read under, if any (see the field doc). + pub(crate) fn query_auth_grant(&self) -> Option<&Arc> { + self.query_auth_grant.as_ref() + } + + /// Whether this split is planned under a grant that filters rows or masks + /// columns. Engines must not publish this split's raw manifest statistics + /// (min/max/null counts describe the data before enforcement), and must not + /// send it across a wire format that cannot carry the grant. + pub fn has_restricted_query_auth_grant(&self) -> bool { + self.query_auth_grant + .as_ref() + .is_some_and(|g| g.has_server_restrictions()) + } + + /// Stamp the query-auth grant this split must be read under. Called by scan + /// planning and write-path authorizers on every emitted split. + pub(crate) fn with_query_auth_grant(mut self, grant: Option>) -> Self { + self.query_auth_grant = grant; + self + } + /// Returns the deletion file for the data file at the given index, if any. `None` at that index means no deletion file. pub fn deletion_file_for_data_file_index(&self, index: usize) -> Option<&DeletionFile> { self.data_deletion_files @@ -652,7 +724,25 @@ impl DataSplit { /// Serialize the DataSplit fields to Java `DataSplit#serialize` (version 8) binary. /// Byte-compatible with `compatibility/datasplit-v8`. Row ranges are not part of the v8 /// format; `serialize_split_v1` wraps a row-range split as an `IndexedSplit` instead. + /// Fail closed when this split carries a restricted query-auth grant and is + /// about to leave the process: no split wire format (native bytes, the Java + /// `SplitSerializer` frame, or serde) carries the grant, so the receiver + /// would read the data raw. `what` names the attempted operation. + fn ensure_no_restricted_grant(&self, what: &str) -> crate::Result<()> { + match &self.query_auth_grant { + Some(grant) if grant.has_server_restrictions() => Err(crate::Error::Unsupported { + message: format!( + "cannot {what} a split planned under a query-auth row filter / column \ + masking grant: the grant cannot cross the wire, so the reader would see \ + unfiltered data" + ), + }), + _ => Ok(()), + } + } + pub fn serialize(&self) -> crate::Result> { + self.ensure_no_restricted_grant("serialize")?; let mut out = Vec::new(); out.extend_from_slice(&SPLIT_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_VERSION.to_be_bytes()); @@ -803,6 +893,10 @@ impl DataSplit { /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. Byte-compatible with /// `compatibility/split-v1-data` / `split-v1-indexed`. pub fn serialize_split_v1(&self) -> crate::Result> { + // The wire format has no place for the query-auth grant (Java carries it + // out of band in `QueryAuthSplit`), so a restricted split would cross the + // boundary as a plain split and be read raw. Fail closed instead. + self.ensure_no_restricted_grant("serialize")?; let mut out = Vec::new(); out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); @@ -1234,6 +1328,7 @@ impl DataSplitBuilder { data_deletion_files: self.data_deletion_files, row_ranges: self.row_ranges, raw_convertible: self.raw_convertible, + query_auth_grant: None, }) } } @@ -1252,15 +1347,45 @@ impl Default for DataSplitBuilder { #[derive(Debug)] pub struct Plan { splits: Vec, + /// False when a residual pass (e.g. a query-auth row filter) drops rows + /// after the scan, so split row counts overcount the read output. + row_counts_exact: bool, } impl Plan { pub fn new(splits: Vec) -> Self { - Self { splits } + Self { + splits, + row_counts_exact: true, + } + } + pub(crate) fn with_inexact_row_counts(mut self) -> Self { + self.row_counts_exact = false; + self + } + + /// Stamp the per-user query-auth grant onto every split so + /// [`crate::table::TableRead::to_arrow`] enforces exactly the grant this + /// plan fetched (see [`DataSplit::with_query_auth_grant`]). A no-op when + /// `grant` is `None` (non-query-auth table): splits keep their empty grant + /// and read raw. + pub(crate) fn stamp_query_auth_grant(mut self, grant: Option>) -> Self { + if grant.is_some() { + self.splits = self + .splits + .into_iter() + .map(|s| s.with_query_auth_grant(grant.clone())) + .collect(); + } + self } pub fn splits(&self) -> &[DataSplit] { &self.splits } + /// Whether split row counts exactly reflect the rows a read will produce. + pub fn row_counts_exact(&self) -> bool { + self.row_counts_exact + } } #[cfg(test)] @@ -1306,6 +1431,41 @@ mod tests { .unwrap() } + #[test] + fn test_restricted_grant_split_cannot_be_serialized() { + use crate::spec::{DataType, IntType, PredicateBuilder}; + use crate::table::query_auth::QueryAuthGrant; + + let fields = vec![crate::spec::DataField::new( + 0, + "id".to_string(), + DataType::Int(IntType::new()), + )]; + let filter = PredicateBuilder::new(&fields) + .greater_than("id", crate::spec::Datum::Int(1)) + .unwrap(); + let restricted = Arc::new(QueryAuthGrant::new(vec![filter], Vec::new(), None)); + let plain = split(vec![file("a", 1, None)], true); + let guarded = plain.clone().with_query_auth_grant(Some(restricted)); + + // No wire format carries the grant, so every serialization path must + // refuse rather than emit a plain split the receiver would read raw. + assert!(guarded.serialize().is_err(), "native bytes"); + assert!(guarded.serialize_split_v1().is_err(), "Java frame"); + assert!( + serde_json::to_vec(&guarded).is_err(), + "serde (pickle, JSON)" + ); + + // Unstamped and unrestricted splits still serialize. + assert!(plain.serialize().is_ok()); + assert!(serde_json::to_vec(&plain).is_ok()); + let unrestricted = plain + .clone() + .with_query_auth_grant(Some(Arc::new(QueryAuthGrant::default()))); + assert!(serde_json::to_vec(&unrestricted).is_ok()); + } + #[test] fn data_split_serde_json_round_trip() { let split = DataSplit::builder() diff --git a/crates/paimon/src/table/table_commit.rs b/crates/paimon/src/table/table_commit.rs index 809e45d9..198170de 100644 --- a/crates/paimon/src/table/table_commit.rs +++ b/crates/paimon/src/table/table_commit.rs @@ -156,6 +156,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; if commit_messages.is_empty() { return Ok(()); @@ -199,6 +200,7 @@ impl TableCommit { commit_identifier: i64, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; if commit_messages.is_empty() { return Ok(()); @@ -270,6 +272,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; if commit_messages.is_empty() && static_partitions.is_none() { return Ok(()); @@ -519,6 +522,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; if partitions.is_empty() { return Ok(()); @@ -566,6 +570,7 @@ impl TableCommit { commit_identifier: i64, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; if partitions.is_empty() { return Err(crate::Error::DataInvalid { @@ -597,6 +602,7 @@ impl TableCommit { filter_committed: bool, ) -> Result<()> { self.table.ensure_not_branch_reference_for_write()?; + self.table.authorize_unrestricted_write().await?; self.try_commit( CommitEntriesPlan::Overwrite { @@ -621,6 +627,10 @@ impl TableCommit { /// files or storage errors are ignored so abort cleanup never masks the /// original write failure. pub async fn abort(&self, commit_messages: &[CommitMessage]) -> Result<()> { + // No query-auth gate: abort only deletes files the caller just wrote and + // publishes nothing. A restricted user's commit is rejected *after* + // `prepare_commit` wrote them, so requiring a grant here would strand + // those files instead of cleaning them up. self.table.ensure_not_branch_reference_for_write()?; for message in commit_messages { diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 6c57f6b6..f81915ac 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -21,7 +21,7 @@ use super::format_table_read::FormatTableRead; use super::incremental_scan::{IncrementalPlan, IncrementalScanMode, IncrementalSplit}; use super::kv_file_reader::{KeyValueFileReader, KeyValueReadConfig}; use super::read_builder::split_scan_predicates; -use super::{ArrowRecordBatchStream, Table}; +use super::{query_auth, ArrowRecordBatchStream, Table}; use crate::arrow::build_target_arrow_schema; use crate::spec::{ BigIntType, CoreOptions, DataField, DataType, MergeEngine, Predicate, TinyIntType, @@ -102,6 +102,15 @@ impl<'a> TableRead<'a> { } } + /// A read-level row limit that must be applied after materialization + /// (format tables); Paimon reads push their limit to scan planning. + fn read_limit(&self) -> Option { + match &self.0 { + TableReadKind::Paimon(_) => None, + TableReadKind::Format(read) => read.limit(), + } + } + /// Set a filter predicate. pub fn with_filter(self, filter: Predicate) -> Self { match self.0 { @@ -127,7 +136,77 @@ impl<'a> TableRead<'a> { } /// Returns an [`ArrowRecordBatchStream`]. + /// + /// Query-auth is enforced here off the grant stamped on the splits by scan + /// planning (or a write-path authorizer), never a shared slot on `Table`, so + /// the grant is exactly the one this query planned and cannot leak from a + /// concurrent query or a write rewrite. A restricted grant is applied on the + /// output stream; an unrestricted grant (or a non-query-auth table) reads + /// raw; and a `query-auth.enabled` table whose splits carry no grant means + /// an unauthorized read path — fail closed. pub fn to_arrow(&self, data_splits: &[DataSplit]) -> crate::Result { + match self.resolve_split_grant(data_splits)? { + Some(grant) if !grant.is_unrestricted() => { + self.to_arrow_auth_enforced(data_splits, grant) + } + _ => self.to_arrow_dispatch(data_splits), + } + } + + /// The single query-auth grant every split in `data_splits` was planned + /// under, or `None` when the table is not `query-auth.enabled`. + /// + /// Splits from different plans must not be mixed: taking any one grant and + /// applying it to the whole slice would let a split carrying a permissive + /// grant relax the read of splits planned under a stricter one, so a + /// disagreement fails closed. A `query-auth.enabled` table whose splits + /// carry no grant means an unauthorized read path — also fail closed. + fn resolve_split_grant( + &self, + data_splits: &[DataSplit], + ) -> crate::Result>> { + let mut grant: Option<&Arc> = None; + let mut saw_ungranted = false; + for split in data_splits { + match (split.query_auth_grant(), &grant) { + (Some(found), None) => grant = Some(found), + (Some(found), Some(seen)) if found != *seen => { + return Err(crate::Error::Unsupported { + message: "reading splits planned under different query-auth grants is \ + not supported; re-plan the scan" + .to_string(), + }); + } + (Some(_), Some(_)) => {} + // A grant found later must not retroactively authorize an + // earlier grant-less split, so record it and check after the + // loop — matching on what was seen so far is order-dependent. + (None, _) => saw_ungranted = true, + } + } + if saw_ungranted && grant.is_some() { + return Err(crate::Error::Unsupported { + message: "a query-auth split was mixed with an unauthorized split; \ + re-plan the scan" + .to_string(), + }); + } + match grant { + Some(grant) => Ok(Some(Arc::clone(grant))), + // Nothing to read: an empty table or a fully pruned scan produces no + // rows, so there is nothing to leak (an authorized plan legitimately + // has no split to stamp). + None if data_splits.is_empty() => Ok(None), + // Real splits with no grant: either not a query-auth table, or an + // unauthorized read path — fail closed. + None => self.table().ensure_read_without_grant().map(|()| None), + } + } + + fn to_arrow_dispatch( + &self, + data_splits: &[DataSplit], + ) -> crate::Result { match &self.0 { TableReadKind::Paimon(read) => read.to_arrow(data_splits), TableReadKind::Format(read) => read.to_arrow(data_splits), @@ -142,6 +221,7 @@ impl<'a> TableRead<'a> { &self, plan: &IncrementalPlan, ) -> crate::Result { + self.ensure_incremental_plan_authorized(plan)?; match &self.0 { TableReadKind::Paimon(read) => read.to_incremental_arrow(plan), TableReadKind::Format(_) => Err(crate::Error::Unsupported { @@ -150,6 +230,22 @@ impl<'a> TableRead<'a> { } } + /// Incremental and audit-log reads consume their splits inside + /// `PaimonTableRead`, which cannot apply the row filter / masking pass, so a + /// restricted grant must fail closed here rather than return raw rows. An + /// unrestricted grant (or a non-query-auth table) proceeds. + fn ensure_incremental_plan_authorized(&self, plan: &IncrementalPlan) -> crate::Result<()> { + let splits = plan.data_splits(); + match self.resolve_split_grant(&splits)? { + Some(grant) if grant.has_server_restrictions() => Err(crate::Error::Unsupported { + message: "reading a query-auth row filter / column masking grant on an \ + incremental or audit-log scan is not supported" + .to_string(), + }), + _ => Ok(()), + } + } + /// Returns an audit-log [`ArrowRecordBatchStream`] for an incremental plan. /// /// Output schema is `rowkind` (+ optional `_SEQUENCE_NUMBER`) followed by @@ -160,6 +256,7 @@ impl<'a> TableRead<'a> { &self, plan: &IncrementalPlan, ) -> crate::Result { + self.ensure_incremental_plan_authorized(plan)?; match &self.0 { TableReadKind::Paimon(read) => read.to_audit_log_arrow(plan), TableReadKind::Format(_) => Err(crate::Error::Unsupported { @@ -167,6 +264,171 @@ impl<'a> TableRead<'a> { }), } } + + /// Read with the query-auth grant applied exactly: read the union of the + /// projection, the filter columns, and the mask inputs; per batch, drop + /// non-matching rows (on raw values, like Java), overwrite masked columns, + /// then project back to the requested columns. + fn to_arrow_auth_enforced( + &self, + data_splits: &[DataSplit], + grant: std::sync::Arc, + ) -> crate::Result { + use futures::StreamExt; + + let table = self.table(); + // The grant's filter and mask indices are POSITIONAL in the schema they + // were parsed against, so enforcing them on another schema would bind + // them to different columns. Refuse a grant issued for a different one. + if !grant.matches_schema(table.schema().id()) { + return Err(crate::Error::Unsupported { + message: "a query-auth grant issued for a different table schema cannot be \ + enforced here; re-plan the scan" + .to_string(), + }); + } + let schema_fields = table.schema().fields().to_vec(); + + // A caller predicate on a masked column would leak raw values through + // row selection (an oracle); refuse such reads. + let masked: std::collections::HashSet = + grant.masks().iter().map(|m| m.column).collect(); + let mut caller_referenced = std::collections::HashSet::new(); + self.data_predicates() + .iter() + .for_each(|p| p.collect_leaf_field_indices(&mut caller_referenced)); + if let Some(index) = caller_referenced.intersection(&masked).next() { + return Err(query_auth::masked_filter_error(&schema_fields, *index)); + } + + // Every projected field must be a canonical (id, name) pair from the + // table schema. Authorization and mask selection resolve columns by id, + // but the physical read resolves them by name, so a hand-built read type + // pairing an authorized id with another column's name would read that + // other column and skip its mask. `with_read_type` / `TableRead::new` + // are public, so this is reachable — fail closed. + for field in self.read_type() { + let by_id = schema_fields.iter().find(|s| s.id() == field.id()); + let by_name = schema_fields.iter().find(|s| s.name() == field.name()); + match (by_id, by_name) { + // Canonical: the same schema field on both lookups. + (Some(s), Some(n)) if s.name() == field.name() && n.id() == field.id() => {} + // Neither an id nor a name of this schema: a system column (row + // id, etc.), which carries no grant scope — leave it to the reader. + (None, None) => {} + // Any mix — a known id under another column's name, or an + // unknown id borrowing a real column's name — would be scoped and + // masked by id while the physical read resolves it by name. + _ => { + return Err(crate::Error::Unsupported { + message: format!( + "query-auth read type field #{} `{}` is not a column of this table", + field.id(), + field.name() + ), + }); + } + } + } + + // The grant is scoped to the columns authorized for this read; a wider + // projection or a predicate on an un-approved column must re-authorize. + // (The builder/scan paths also check this before pruning; this covers a + // directly-constructed `TableRead`.) + let projected_indices = self + .read_type() + .iter() + .filter_map(|f| schema_fields.iter().position(|s| s.id() == f.id())); + if let Some(index) = + grant.first_unauthorized(projected_indices.chain(caller_referenced.iter().copied())) + { + return Err(query_auth::unauthorized_column_error(&schema_fields, index)); + } + + // Only masks whose target is caller-projected matter (others are + // projected away); keeping just those avoids masking a target that was + // added to the physical read solely because a filter references it. + let caller_ids: std::collections::HashSet = + self.read_type().iter().map(|f| f.id()).collect(); + let masks: Vec = grant + .masks() + .iter() + .filter(|m| { + schema_fields + .get(m.column) + .is_some_and(|t| caller_ids.contains(&t.id())) + }) + .cloned() + .collect(); + + // Widen the physical read with filter columns and the applied masks' + // inputs, so both are always available to the in-memory pass. + let mut referenced = std::collections::HashSet::new(); + grant + .filters() + .iter() + .for_each(|f| f.collect_leaf_field_indices(&mut referenced)); + masks + .iter() + .for_each(|m| m.transform.collect_field_indices(&mut referenced)); + let mut physical = self.read_type().to_vec(); + for index in referenced { + let field = schema_fields + .get(index) + .ok_or_else(|| crate::Error::Unsupported { + message: format!("query-auth grant references unknown field #{index}"), + })?; + if !physical.iter().any(|f| f.id() == field.id()) { + physical.push(field.clone()); + } + } + + let projected_columns = self.read_type().len(); + let filters = grant.filters().to_vec(); + // The inner read must NOT apply the caller's limit: it would cap rows + // before the auth filter. Read everything, then truncate the output. + let caller_limit = self.read_limit(); + let inner = TableRead::new(table, physical.clone(), self.data_predicates().to_vec()); + let stream = inner.to_arrow_dispatch(data_splits)?.map(move |batch| { + let batch = batch?; + let filtered = + query_auth::strict_filter_batch(&batch, &filters, &schema_fields, &physical)?; + let masked = query_auth::mask_batch(&filtered, &masks, &schema_fields, &physical)?; + masked + .project(&(0..projected_columns).collect::>()) + .map_err(|e| crate::Error::DataInvalid { + message: format!("failed to re-project query-auth batch: {e}"), + source: Some(Box::new(e)), + }) + }); + match caller_limit { + None => Ok(Box::pin(stream)), + // `unfold` stops as soon as the cap is reached (state `emitted >= + // limit`) without polling the inner stream again, so a read error in + // a later batch can't surface after the limit is satisfied. + Some(limit) => Ok(Box::pin(futures::stream::unfold( + (Box::pin(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 remaining = limit - emitted; + let batch = if batch.num_rows() > remaining { + batch.slice(0, remaining) + } else { + batch + }; + let emitted = emitted + batch.num_rows(); + Some((Ok(batch), (inner, emitted))) + } + } + }, + ))), + } + } } #[derive(Debug, Clone)] @@ -361,12 +623,12 @@ impl<'a> PaimonTableRead<'a> { })) } - /// Returns an [`ArrowRecordBatchStream`]. + /// Returns an [`ArrowRecordBatchStream`]. Query-auth (fail-closed + row + /// filter + masking) is enforced by the outer [`TableRead::to_arrow`] off + /// the grant stamped on the splits. 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(); - // Fail closed for a direct `TableRead` (bypassing `ReadBuilder::new_read`). - core_options.ensure_read_authorized()?; let merge_engine = core_options.merge_engine()?; // Route supported PK merge engines through the split-aware reader. @@ -723,12 +985,97 @@ mod tests { // Bypass `ReadBuilder` by constructing `TableRead` directly; the `to_arrow` guard // still fails closed. let read = TableRead::new(&table, table.schema.fields().to_vec(), Vec::new()); + // Real splits with no stamped grant: the read would return raw rows. + let ungranted = split(vec![file("a", 5, Some(0))], true); assert!( matches!( - read.to_arrow(&[]), + read.to_arrow(&[ungranted]), Err(crate::Error::Unsupported { ref message }) if message.contains("query-auth.enabled") ), "directly-constructed read of a query-auth.enabled table must fail closed" ); + // An empty slice reads no rows, so it is allowed (an authorized plan may + // legitimately produce no splits). + assert!(read.to_arrow(&[]).is_ok()); + } + + #[test] + fn test_noncanonical_read_type_fails_closed() { + use std::collections::HashSet; + + // Authorization and mask selection resolve by field id, but the physical + // read resolves by name. A read type pairing an authorized id with + // another column's name would read that other column and skip its mask. + let table = query_auth_table(); + let schema_field = table.schema.fields()[0].clone(); + let forged = crate::spec::DataField::new( + schema_field.id(), + "not_the_real_name".to_string(), + schema_field.data_type().clone(), + ); + let grant = Arc::new(query_auth::QueryAuthGrant::new( + Vec::new(), + Vec::new(), + Some(HashSet::from([0])), + )); + let granted = split(vec![file("a", 5, Some(0))], true).with_query_auth_grant(Some(grant)); + + // Direction A: a known id carrying another column's name. + let read = TableRead::new(&table, vec![forged], Vec::new()); + let Err(err) = read.to_arrow(std::slice::from_ref(&granted)) else { + panic!("a read type with a mismatched id/name pair must fail closed"); + }; + assert!( + err.to_string().contains("not a column of this table"), + "got: {err}" + ); + + // Direction B: an UNKNOWN id borrowing a real column's name. The + // physical read resolves by name, so this would read the real column + // while scoping and mask selection (both by id) see an unrelated field. + let forged_id = crate::spec::DataField::new( + 9999, + schema_field.name().to_string(), + schema_field.data_type().clone(), + ); + let read = TableRead::new(&table, vec![forged_id], Vec::new()); + let Err(err) = read.to_arrow(&[granted]) else { + panic!("an unknown id borrowing a real column name must fail closed"); + }; + assert!( + err.to_string().contains("not a column of this table"), + "got: {err}" + ); + } + + #[test] + fn test_mixed_granted_and_ungranted_splits_fail_closed_in_both_orders() { + use std::collections::HashSet; + + let table = query_auth_table(); + let read = TableRead::new(&table, table.schema.fields().to_vec(), Vec::new()); + let grant = Arc::new(query_auth::QueryAuthGrant::new( + Vec::new(), + Vec::new(), + Some(HashSet::from([0])), + )); + let granted = + split(vec![file("a", 5, Some(0))], true).with_query_auth_grant(Some(grant.clone())); + let ungranted = split(vec![file("b", 5, Some(0))], true); + + // Order must not decide the verdict: a grant found later must never + // retroactively authorize an earlier grant-less split. + for slice in [ + vec![granted.clone(), ungranted.clone()], + vec![ungranted, granted], + ] { + let Err(err) = read.to_arrow(&slice) else { + panic!("mixing a granted and an ungranted split must fail closed"); + }; + assert!( + err.to_string().contains("mixed with an unauthorized split"), + "got: {err}" + ); + } } } diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index c8c0f199..597c39a3 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -44,6 +44,7 @@ use crate::table::bin_pack::split_for_batch; use crate::table::merge_tree_split_generator::{ merge_tree_split_for_batch, KeyComparator, SplitGroup, }; +use crate::table::query_auth::QueryAuthGrant; use crate::table::schema_manager::SchemaManager; use crate::table::source::{ any_range_overlaps_file, intersect_ranges_with_file, merge_row_ranges, DataSplit, @@ -941,6 +942,21 @@ impl<'a> TableScan<'a> { } } + pub(super) fn with_query_auth_scope( + self, + filter_columns: HashSet, + projected: Option>, + ) -> Self { + match self.0 { + TableScanKind::Paimon(scan) => Self(TableScanKind::Paimon( + scan.with_query_auth_scope(filter_columns, projected), + )), + TableScanKind::Format(scan) => Self(TableScanKind::Format( + scan.with_query_auth_scope(filter_columns, projected), + )), + } + } + pub async fn plan(&self) -> crate::Result { match &self.0 { TableScanKind::Paimon(scan) => scan.plan().await, @@ -991,7 +1007,7 @@ impl<'a> TableScan<'a> { fn apply_limit_pushdown(&self, splits: Vec) -> Vec { match &self.0 { TableScanKind::Paimon(scan) => scan.apply_limit_pushdown(splits), - TableScanKind::Format(scan) => scan.apply_limit_pushdown(splits), + TableScanKind::Format(scan) => scan.apply_limit_pushdown(splits, false), } } } @@ -1014,6 +1030,10 @@ struct PaimonTableScan<'a> { /// the complete file set. Normal read scans leave this as `false`. scan_all_files: bool, projected_read_field_ids: Option>, + /// Filter/projection columns for the query-auth scope check, evaluated + /// against the live grant at plan time (see `ensure_query_auth_allowed`). + query_auth_filter_columns: HashSet, + query_auth_projected: Option>, } impl<'a> PaimonTableScan<'a> { @@ -1034,6 +1054,8 @@ impl<'a> PaimonTableScan<'a> { row_ranges, scan_all_files: false, projected_read_field_ids: None, + query_auth_filter_columns: HashSet::new(), + query_auth_projected: None, } } @@ -1068,6 +1090,16 @@ impl<'a> PaimonTableScan<'a> { self } + pub(super) fn with_query_auth_scope( + mut self, + filter_columns: HashSet, + projected: Option>, + ) -> Self { + self.query_auth_filter_columns = filter_columns; + self.query_auth_projected = projected; + self + } + /// Plan the full scan: resolve snapshot (via options or latest), then read manifests and build DataSplits. /// /// Time travel is resolved from table options: @@ -1084,27 +1116,40 @@ impl<'a> PaimonTableScan<'a> { /// for `scan.version`; the strict selectors mirror Java's typed /// `scan.snapshot-id` / `scan.tag-name` handling. pub async fn plan(&self) -> crate::Result { - self.ensure_query_auth_allowed()?; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; + let grant = self.ensure_query_auth_allowed().await?; + let has_row_filter = grant.as_deref().is_some_and(|g| g.has_row_filter()); + let data_evolution_read_field_ids = self.auth_widened_read_field_ids(grant.as_deref()); let snapshot = match self.resolve_snapshot().await? { Some(snapshot) => snapshot, - None => return Ok(Plan::new(Vec::new())), + None => return Ok(self.finalize_plan(Plan::new(Vec::new()), grant.as_ref())), }; - self.plan_snapshot(snapshot, data_evolution_read_field_ids.as_ref(), None) - .await + self.plan_snapshot( + snapshot, + data_evolution_read_field_ids.as_ref(), + None, + has_row_filter, + ) + .await + .map(|plan| self.finalize_plan(plan, grant.as_ref())) } /// Plan the full scan and return metadata-pruning trace counters. pub async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { - self.ensure_query_auth_allowed()?; + let grant = self.ensure_query_auth_allowed().await?; + let has_row_filter = grant.as_deref().is_some_and(|g| g.has_row_filter()); let mut trace = ScanTrace { limit: self.limit, ..Default::default() }; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; + let data_evolution_read_field_ids = self.auth_widened_read_field_ids(grant.as_deref()); let snapshot = match self.resolve_snapshot().await? { Some(snapshot) => snapshot, - None => return Ok((Plan::new(Vec::new()), trace)), + None => { + return Ok(( + self.finalize_plan(Plan::new(Vec::new()), grant.as_ref()), + trace, + )) + } }; trace.snapshot_id = Some(snapshot.id()); let plan = self @@ -1112,20 +1157,70 @@ impl<'a> PaimonTableScan<'a> { snapshot, data_evolution_read_field_ids.as_ref(), Some(&mut trace), + has_row_filter, ) .await?; - Ok((plan, trace)) + Ok((self.finalize_plan(plan, grant.as_ref()), trace)) + } + + /// Stamp the grant onto every split (so `TableRead::to_arrow` enforces + /// exactly this plan's grant) and mark row counts inexact when it carries a + /// row filter (dropped as a residual pass inside `TableRead`). + fn finalize_plan(&self, plan: Plan, grant: Option<&Arc>) -> Plan { + // Any restricted grant invalidates the plan's statistics: a row filter + // drops rows, and masking rewrites column values (a NULL mask makes the + // column entirely null), so a statistics-only `COUNT` would report raw + // counts and bypass enforcement. + let restricted = grant.is_some_and(|g| g.has_server_restrictions()); + let plan = plan.stamp_query_auth_grant(grant.cloned()); + if restricted { + plan.with_inexact_row_counts() + } else { + plan + } } /// Fail closed for a `query-auth.enabled` table: scan planning — including /// `with_scan_all_files`, which read-facing system tables like `files` use — /// exposes file paths, row counts, and stats the client can't authorize. - fn ensure_query_auth_allowed(&self) -> crate::Result<()> { - CoreOptions::new(self.table.schema().options()).ensure_read_authorized() - } - - fn projected_read_field_ids(&self) -> crate::Result>> { - Ok(self.projected_read_field_ids.clone()) + /// Returns the fetched grant (`None` = not a query-auth table) so the caller + /// can widen the projection and stamp the splits with it. + async fn ensure_query_auth_allowed(&self) -> crate::Result>> { + // Fetch/refresh the grant at plan time (Java parity), then guard + // against pruning on masked or out-of-scope columns. + let select = self.query_auth_projected.as_ref().map(|projected| { + projected + .iter() + .copied() + .chain(self.query_auth_filter_columns.iter().copied()) + .collect::>() + }); + let grant = self + .table + .verify_query_auth_for_read(select.as_ref()) + .await?; + if let Some(grant) = &grant { + crate::table::query_auth::scope_check( + grant, + self.table.schema().fields(), + &self.query_auth_filter_columns, + self.query_auth_projected.clone(), + )?; + } + Ok(grant) + } + + /// The projected field ids for data-evolution column-slice pruning, widened + /// with the grant's row-filter / mask-input columns so pruning cannot drop a + /// file holding a column the grant needs (an omitted column would read as + /// null and wrongly satisfy `IS_NULL`). Computed here — with the grant in + /// hand — rather than at `new_scan` time, where the grant does not yet exist. + fn auth_widened_read_field_ids(&self, grant: Option<&QueryAuthGrant>) -> Option> { + let mut ids = self.projected_read_field_ids.clone(); + if let (Some(set), Some(grant)) = (ids.as_mut(), grant) { + set.extend(grant.read_field_ids(self.table.schema().fields())); + } + ids } async fn resolve_snapshot(&self) -> crate::Result> { @@ -1290,8 +1385,16 @@ impl<'a> PaimonTableScan<'a> { Ok(merged) } - fn can_push_down_limit_hint(&self, row_ranges: Option<&[RowRange]>) -> bool { + fn can_push_down_limit_hint( + &self, + row_ranges: Option<&[RowRange]>, + query_auth_row_filter: bool, + ) -> bool { + // A query-auth row filter is applied as a residual pass at read time, so + // split merged_row_count overcounts; count-based limit pruning would drop + // splits holding later authorized rows. can_push_down_limit_hint_for_scan(&self.data_predicates, row_ranges) + && !query_auth_row_filter } fn global_index_scan_settings( @@ -1452,14 +1555,17 @@ impl<'a> PaimonTableScan<'a> { /// Reuses the same split-building path as a full snapshot plan, but only /// reads the delta manifest list and keeps ADD entries. pub(crate) async fn plan_snapshot_delta(&self, snapshot: &Snapshot) -> crate::Result { - self.ensure_query_auth_allowed()?; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; - self.plan_snapshot_manifest_list( - snapshot, - snapshot.delta_manifest_list(), - data_evolution_read_field_ids.as_ref(), - ) - .await + let grant = self.ensure_query_auth_allowed().await?; + let data_evolution_read_field_ids = self.auth_widened_read_field_ids(grant.as_deref()); + let plan = self + .plan_snapshot_manifest_list( + snapshot, + snapshot.delta_manifest_list(), + data_evolution_read_field_ids.as_ref(), + grant.as_deref().is_some_and(|g| g.has_row_filter()), + ) + .await?; + Ok(self.finalize_plan(plan, grant.as_ref())) } /// Plan data splits from a snapshot's changelog manifest list. @@ -1468,17 +1574,20 @@ impl<'a> PaimonTableScan<'a> { /// reads the changelog manifest list and keeps ADD entries. Snapshots /// without a changelog list yield an empty plan. pub(crate) async fn plan_snapshot_changelog(&self, snapshot: &Snapshot) -> crate::Result { - self.ensure_query_auth_allowed()?; + let grant = self.ensure_query_auth_allowed().await?; let Some(list_name) = snapshot.changelog_manifest_list() else { - return Ok(Plan::new(Vec::new())); + return Ok(self.finalize_plan(Plan::new(Vec::new()), grant.as_ref())); }; - let data_evolution_read_field_ids = self.projected_read_field_ids()?; - self.plan_snapshot_manifest_list( - snapshot, - list_name, - data_evolution_read_field_ids.as_ref(), - ) - .await + let data_evolution_read_field_ids = self.auth_widened_read_field_ids(grant.as_deref()); + let plan = self + .plan_snapshot_manifest_list( + snapshot, + list_name, + data_evolution_read_field_ids.as_ref(), + grant.as_deref().is_some_and(|g| g.has_row_filter()), + ) + .await?; + Ok(self.finalize_plan(plan, grant.as_ref())) } async fn plan_snapshot_manifest_list( @@ -1486,6 +1595,7 @@ impl<'a> PaimonTableScan<'a> { snapshot: &Snapshot, manifest_list_name: &str, data_evolution_read_field_ids: Option<&HashSet>, + query_auth_row_filter: bool, ) -> crate::Result { if matches!(self.limit, Some(0)) { return Ok(Plan::new(Vec::new())); @@ -1531,6 +1641,7 @@ impl<'a> PaimonTableScan<'a> { index_entries, effective_row_ranges, None, + query_auth_row_filter, ) .await } @@ -1659,6 +1770,7 @@ impl<'a> PaimonTableScan<'a> { snapshot: Snapshot, data_evolution_read_field_ids: Option<&HashSet>, mut trace: Option<&mut ScanTrace>, + query_auth_row_filter: bool, ) -> crate::Result { if matches!(self.limit, Some(0)) { if let Some(trace) = trace { @@ -1714,10 +1826,12 @@ impl<'a> PaimonTableScan<'a> { index_entries, effective_row_ranges, trace, + query_auth_row_filter, ) .await } + #[allow(clippy::too_many_arguments)] async fn plan_snapshot_from_entries( &self, snapshot: Snapshot, @@ -1726,6 +1840,7 @@ impl<'a> PaimonTableScan<'a> { index_entries: Option>, effective_row_ranges: Option>, mut trace: Option<&mut ScanTrace>, + query_auth_row_filter: bool, ) -> crate::Result { let table_path = self.table.location(); let table_schema_id = self.table.schema().id(); @@ -1853,7 +1968,8 @@ impl<'a> PaimonTableScan<'a> { .map(|entries| build_deletion_files_map(entries, base_path)); let mut data_file_field_ids_cache = DataFileFieldIdsCache::new(); - let can_push_down_limit = self.can_push_down_limit_hint(effective_row_ranges.as_deref()); + let can_push_down_limit = + self.can_push_down_limit_hint(effective_row_ranges.as_deref(), query_auth_row_filter); let mut limit_accumulator = match self.limit { Some(limit) if limit > 0 && can_push_down_limit => { Some(LimitPushdownAccumulator::new(limit)) diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 2c60aa9a..69a68c08 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -200,8 +200,10 @@ impl<'a> VectorSearchBuilder<'a> { pub async fn execute_scored(&self) -> crate::Result { // Fail closed: returns data-derived row ranges outside `TableScan`/`TableRead`. + // Strict: search results bypass the query-auth row filter, so only a + // fully unrestricted grant may search. + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -275,9 +277,11 @@ impl<'a> VectorSearchBuilder<'a> { /// [`with_projection`](Self::with_projection)) plus `__paimon_search_score`; /// `_ROW_ID` and `_PKEY_VECTOR_POSITION` are always hidden. pub async fn execute_read(&self) -> crate::Result { - // Fail closed: returns data outside `TableScan`/`TableRead`. + // Fail closed: materializes rows outside `TableScan`/`TableRead`, so it + // cannot apply the row filter / masking — only a fully unrestricted + // grant may run it (same gate as the scored entry points). + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -1053,12 +1057,14 @@ impl<'a> BatchVectorSearchBuilder<'a> { } pub async fn execute(&self) -> crate::Result> { - // Fail closed: like `execute_read` and the single-query builder, this - // returns data-derived row ids/scores outside `TableScan`/`TableRead`, - // so it must refuse a `query-auth.enabled` table before any fast path - // (an empty snapshot would otherwise return empty results and bypass it). + // Strict, like the scored path: batch vector search reads index files + // raw and returns top-k membership/ordering/scores over masked or + // filter-hidden rows — a ranking oracle it cannot enforce the row filter + // on, so only a fully unrestricted grant may run it. Checked before any + // validation or fast path (an empty snapshot would otherwise return + // empty results and bypass it); also covers the DataFusion lateral path. + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -1172,9 +1178,11 @@ impl<'a> BatchVectorSearchBuilder<'a> { /// scored global row-ids, not materialized rows, so callers use /// [`execute`](Self::execute) instead. pub async fn execute_read(&self) -> crate::Result> { - // Fail closed: returns data outside `TableScan`/`TableRead`. + // Fail closed: materializes rows outside `TableScan`/`TableRead`, so it + // cannot apply the row filter / masking — only a fully unrestricted + // grant may run it (same gate as the scored entry points). + self.table.authorize_unrestricted_read().await?; let core = CoreOptions::new(self.table.schema().options()); - core.ensure_read_authorized()?; let vector_column = self.vector_column .as_deref() @@ -3532,11 +3540,11 @@ mod tests { #[tokio::test] async fn test_batch_execute_fails_closed_when_query_auth_enabled() { - // The batch scored entry returns data-derived row ids/scores outside - // `TableScan`/`TableRead`, so it must fail closed under - // `query-auth.enabled` exactly like the single-query builder. Its config - // is otherwise valid, so without the guard the empty-snapshot fast path - // would return empty results and silently bypass authorization. + // The batch builder reads index files raw (never through plan/to_arrow), + // so it must gate query-auth itself — a top-k ranking oracle over + // masked/filter-hidden rows otherwise. The config here is valid, so + // without the guard the empty-snapshot fast path would return empty + // results and silently bypass authorization. let table = crate::table::query_auth_table(); let err = table .new_batch_vector_search_builder() diff --git a/crates/paimon/src/table/vindex_index_build_builder.rs b/crates/paimon/src/table/vindex_index_build_builder.rs index 11579c6e..3d94026e 100644 --- a/crates/paimon/src/table/vindex_index_build_builder.rs +++ b/crates/paimon/src/table/vindex_index_build_builder.rs @@ -552,6 +552,10 @@ async fn extract_vectors( index_column: &str, dimension: i32, ) -> Result> { + // Index building reads raw values, so it requires an unrestricted grant (a + // restricted one would index a filtered/masked view); stamp it so the read + // is authorized. Mirrors the B-tree index builder. + let build_grant = table.authorize_unrestricted_read().await?; let split = DataSplitBuilder::new() .with_snapshot(shard.snapshot_id) .with_partition(shard.partition.clone()) @@ -563,7 +567,8 @@ async fn extract_vectors( shard.row_range_start, shard.row_range_end, )]) - .build()?; + .build()? + .with_query_auth_grant(build_grant); let mut read_builder = table.new_read_builder(); read_builder.with_projection(&[index_column, ROW_ID_FIELD_NAME])?; diff --git a/crates/paimon/tests/mock_server.rs b/crates/paimon/tests/mock_server.rs index 6c8a0048..778247b3 100644 --- a/crates/paimon/tests/mock_server.rs +++ b/crates/paimon/tests/mock_server.rs @@ -34,8 +34,8 @@ use std::sync::{Arc, Mutex}; use tokio::task::JoinHandle; use paimon::api::{ - AlterDatabaseRequest, AlterTableRequest, AuditRESTResponse, ConfigResponse, - CreateFunctionRequest, CreateViewRequest, ErrorResponse, GetDatabaseResponse, + AlterDatabaseRequest, AlterTableRequest, AuditRESTResponse, AuthTableQueryResponse, + ConfigResponse, CreateFunctionRequest, CreateViewRequest, ErrorResponse, GetDatabaseResponse, GetFunctionResponse, GetTableResponse, GetViewResponse, ListDatabasesResponse, ListFunctionsResponse, ListTablesResponse, ListViewsResponse, RenameTableRequest, ResourcePaths, @@ -53,6 +53,8 @@ struct MockState { list_page_size: Option, no_permission_databases: HashSet, no_permission_tables: HashSet, + /// Per-table auth response for `POST .../tables/{table}/auth`; absent = unrestricted. + auth_responses: HashMap, /// ECS metadata role name (for token loader testing) ecs_role_name: Option, /// ECS metadata token (for token loader testing) @@ -705,6 +707,20 @@ impl RESTServer { (StatusCode::NOT_FOUND, Json(err)).into_response() } + /// Handle POST /databases/:db/tables/:table/auth - per-user query-auth check. + pub async fn auth_table_query( + Path((db, table)): Path<(String, String)>, + Extension(state): Extension>, + ) -> impl IntoResponse { + let s = state.inner.lock().unwrap(); + let response = s + .auth_responses + .get(&format!("{db}.{table}")) + .cloned() + .unwrap_or_default(); + (StatusCode::OK, Json(response)).into_response() + } + /// Handle DELETE /databases/:db/tables/:table - drop a table. pub async fn drop_table( Path((db, table)): Path<(String, String)>, @@ -967,6 +983,13 @@ impl RESTServer { ); } + /// Set the auth response returned for `POST .../tables/{table}/auth`. + pub fn set_auth_response(&self, database: &str, table: &str, response: AuthTableQueryResponse) { + let mut s = self.inner.lock().unwrap(); + s.auth_responses + .insert(format!("{database}.{table}"), response); + } + /// Add a no-permission table to the server state. pub fn add_no_permission_table(&self, database: &str, table: &str) { let mut s = self.inner.lock().unwrap(); @@ -1105,6 +1128,10 @@ pub async fn start_mock_server( &format!("{prefix}/databases/:db/functions/:function"), get(RESTServer::get_function), ) + .route( + &format!("{prefix}/databases/:db/tables/:table/auth"), + post(RESTServer::auth_table_query), + ) .route( &format!("{prefix}/tables/rename"), post(RESTServer::rename_table), diff --git a/crates/paimon/tests/rest_catalog_test.rs b/crates/paimon/tests/rest_catalog_test.rs index 2000f9e4..5414303f 100644 --- a/crates/paimon/tests/rest_catalog_test.rs +++ b/crates/paimon/tests/rest_catalog_test.rs @@ -27,7 +27,7 @@ use arrow_array::{Array, BinaryArray, Int32Array, Int64Array, RecordBatch, Strin use arrow_schema::{DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema}; use axum::http::StatusCode; use futures::TryStreamExt; -use paimon::api::ConfigResponse; +use paimon::api::{AuthTableQueryResponse, ConfigResponse}; use paimon::catalog::{Catalog, Function, FunctionDefinition, Identifier, RESTCatalog, ViewSchema}; use paimon::common::Options; use paimon::spec::{ @@ -346,6 +346,564 @@ async fn test_catalog_get_table() { assert!(table.is_ok(), "failed to get table: {table:?}"); } +/// The grant is scoped to the columns the plan requested (like Java passing +/// `readType.getFieldNames()`): a wider read against a scoped grant fails +/// closed until it re-plans. +#[tokio::test] +async fn test_query_auth_grant_scoped_to_planned_columns() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = Schema::builder() + .column("id", DataType::BigInt(BigIntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .option("query-auth.enabled", "true") + .build() + .unwrap(); + ctx.server.add_table_with_schema( + "default", + "qa_scope", + schema, + "file:///tmp/test_warehouse/default.db/qa_scope", + ); + let table = ctx + .catalog + .get_table(&Identifier::new("default", "qa_scope")) + .await + .unwrap(); + + // Plan a projection of {id}: the grant is scoped to that column and stamped + // on this plan's splits (not a shared slot on the table). + let mut projected = table.new_read_builder(); + projected.with_projection(&["id"]).unwrap(); + projected.new_scan().plan().await.unwrap(); + + // The {id} plan's scoped grant lives on its own splits, so it cannot leak + // to a separate full-table read (per-query, no shared mutable slot). Reading + // real splits that carry no grant fails closed — covered by the read_builder + // unit tests, which can build such a split directly. + + // A separate full-table read re-authorizes for all columns when it plans. + table.new_read_builder().new_scan().plan().await.unwrap(); +} + +/// Java #8447 baseline: a query-auth row filter disables count-based limit +/// pushdown, so a limited read still reaches authorized rows in later files. +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_row_filter_reads_past_limit_pushdown() { + let tmp = tempfile::tempdir().unwrap(); + let warehouse = format!("file://{}", tmp.path().display()); + + // Write two commits (-> two files) through a plain FileSystemCatalog table. + let mut fs_options = Options::new(); + fs_options.set(CatalogOptions::WAREHOUSE, &warehouse); + let fs_catalog = FileSystemCatalog::new(fs_options).expect("create filesystem catalog"); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let write_schema = Schema::builder() + .column("id", DataType::BigInt(BigIntType::new())) + .option("bucket", "1") + .option("bucket-key", "id") + .build() + .unwrap(); + let identifier = Identifier::new("default", "qa_limit"); + fs_catalog + .create_table(&identifier, write_schema, false) + .await + .unwrap(); + let writer = fs_catalog.get_table(&identifier).await.unwrap(); + let int_batch = |ids: Vec| { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + ArrowDataType::Int64, + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids))]).unwrap() + }; + write_batch(&writer, int_batch(vec![1, 2, 3, 4]), "u1").await; + write_batch(&writer, int_batch(vec![5, 6, 7, 8]), "u2").await; + + // Read the same files through the REST catalog with query-auth enabled and + // a row filter of id >= 6 (all matches live in the SECOND file). + let ctx = setup_catalog(vec!["default"]).await; + let read_schema = Schema::builder() + .column("id", DataType::BigInt(BigIntType::new())) + .option("bucket", "1") + .option("bucket-key", "id") + .option("query-auth.enabled", "true") + .build() + .unwrap(); + ctx.server.add_table_with_schema( + "default", + "qa_limit", + read_schema, + &format!("{warehouse}/default.db/qa_limit"), + ); + ctx.server.set_auth_response( + "default", + "qa_limit", + AuthTableQueryResponse { + filter: Some(vec![ + r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":0,"name":"id","type":"BIGINT"}},"function":"GREATER_OR_EQUAL","literals":[6]}"# + .to_string(), + ]), + column_masking: None, + }, + ); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + + let mut builder = table.new_read_builder(); + builder.with_limit(2); + let plan = builder.new_scan().plan().await.unwrap(); + // The filter runs as a residual pass, so the plan must not cap splits by + // the unfiltered limit and must report its row counts as inexact. + assert!(!plan.row_counts_exact()); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + ids.sort_unstable(); + assert_eq!( + ids, + vec![6, 7, 8], + "authorized rows beyond file 1 must appear" + ); +} + +/// Java #8570 baseline: a cross-column mask (`alias := UPPER(name)`) and a row +/// filter on an unprojected column (`score`) must still enforce when the caller +/// projects neither `name` nor `score` — the read is widened with the grant's +/// columns and every batch of every split is projected back to the caller's +/// columns (no auth-added column may leak from later splits). +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_cross_column_mask_with_narrow_projection() { + let tmp = tempfile::tempdir().unwrap(); + let warehouse = format!("file://{}", tmp.path().display()); + + let mut fs_options = Options::new(); + fs_options.set(CatalogOptions::WAREHOUSE, &warehouse); + let fs_catalog = FileSystemCatalog::new(fs_options).expect("create filesystem catalog"); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let columns = || { + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("alias", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("score", DataType::BigInt(BigIntType::new())) + .option("bucket", "1") + .option("bucket-key", "id") + }; + let identifier = Identifier::new("default", "qa_cross_mask"); + fs_catalog + .create_table(&identifier, columns().build().unwrap(), false) + .await + .unwrap(); + let writer = fs_catalog.get_table(&identifier).await.unwrap(); + let batch = |ids: Vec, names: Vec<&str>, aliases: Vec<&str>, scores: Vec| { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("alias", ArrowDataType::Utf8, true), + ArrowField::new("score", ArrowDataType::Int64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + Arc::new(StringArray::from(aliases)), + Arc::new(Int64Array::from(scores)), + ], + ) + .unwrap() + }; + // Two commits -> two files, so enforcement is exercised across splits. + write_batch( + &writer, + batch(vec![1, 2], vec!["ann", "bob"], vec!["x", "y"], vec![5, 15]), + "u1", + ) + .await; + write_batch( + &writer, + batch(vec![3, 4], vec!["cid", "dan"], vec!["z", "w"], vec![20, 8]), + "u2", + ) + .await; + + let ctx = setup_catalog(vec!["default"]).await; + ctx.server.add_table_with_schema( + "default", + "qa_cross_mask", + columns() + .option("query-auth.enabled", "true") + .build() + .unwrap(), + &format!("{warehouse}/default.db/qa_cross_mask"), + ); + ctx.server.set_auth_response( + "default", + "qa_cross_mask", + AuthTableQueryResponse { + // Row filter on `score` (index 3), which the caller does not project. + filter: Some(vec![ + r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":3,"name":"score","type":"BIGINT"}},"function":"GREATER_OR_EQUAL","literals":[10]}"# + .to_string(), + ]), + // Cross-column mask: `alias` is overwritten from `name` (index 1), + // which the caller does not project either. + column_masking: Some(HashMap::from([( + "alias".to_string(), + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + .to_string(), + )])), + }, + ); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + + let mut builder = table.new_read_builder(); + builder.with_projection(&["id", "alias"]).unwrap(); + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + + let mut rows = Vec::new(); + for b in &batches { + assert_eq!( + b.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["id", "alias"], + "auth-added columns (name, score) must not leak from any split" + ); + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let aliases = b.column(1).as_any().downcast_ref::().unwrap(); + for r in 0..b.num_rows() { + rows.push((ids.value(r), aliases.value(r).to_string())); + } + } + rows.sort_unstable(); + assert_eq!( + rows, + vec![(2, "BOB".to_string()), (3, "CID".to_string())], + "filter must drop score<10 rows in both files and alias must be UPPER(name)" + ); +} + +/// Visible end-to-end demo of query-auth enforcement over a mock REST catalog: +/// the same files are read once with no grant (raw) and once through a +/// `query-auth.enabled` table whose per-user grant applies a row filter +/// (`salary >= 90000`) plus column masking (`name -> UPPER(name)`). Run with: +/// cargo test -p paimon --test rest_catalog_test query_auth_enforcement_demo -- --nocapture +#[cfg(not(windows))] +#[tokio::test] +async fn test_query_auth_enforcement_demo() { + let tmp = tempfile::tempdir().unwrap(); + let warehouse = format!("file://{}", tmp.path().display()); + + // --- Write demo_employees (id, name, salary) via a plain FileSystemCatalog --- + let mut fs_options = Options::new(); + fs_options.set(CatalogOptions::WAREHOUSE, &warehouse); + let fs_catalog = FileSystemCatalog::new(fs_options).expect("create filesystem catalog"); + fs_catalog + .create_database("default", true, HashMap::new()) + .await + .unwrap(); + let columns = || { + Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::new(255).unwrap())) + .column("salary", DataType::BigInt(BigIntType::new())) + .option("bucket", "1") + .option("bucket-key", "id") + }; + let identifier = Identifier::new("default", "demo_employees"); + fs_catalog + .create_table(&identifier, columns().build().unwrap(), false) + .await + .unwrap(); + let writer = fs_catalog.get_table(&identifier).await.unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("salary", ArrowDataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec![ + "alice", "bob", "charlie", "diana", "eve", + ])), + Arc::new(Int64Array::from(vec![120000, 85000, 95000, 70000, 99000])), + ], + ) + .unwrap(); + write_batch(&writer, batch, "u1").await; + + let dump = |label: &str, batches: &[RecordBatch]| { + println!("\n {label}"); + println!(" {:<4} {:<10} {:>8}", "id", "name", "salary"); + for b in batches { + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let names = b.column(1).as_any().downcast_ref::().unwrap(); + let sal = b.column(2).as_any().downcast_ref::().unwrap(); + for r in 0..b.num_rows() { + println!( + " {:<4} {:<10} {:>8}", + ids.value(r), + names.value(r), + sal.value(r) + ); + } + } + }; + + // --- Raw read (no query-auth) --- + let raw: Vec = { + let b = writer.new_read_builder(); + b.new_read() + .unwrap() + .to_arrow(b.new_scan().plan().await.unwrap().splits()) + .unwrap() + .try_collect() + .await + .unwrap() + }; + dump("RAW (no grant): all rows, real names", &raw); + + // --- Enforced read through the REST catalog with a per-user grant --- + let ctx = setup_catalog(vec!["default"]).await; + ctx.server.add_table_with_schema( + "default", + "demo_employees", + columns() + .option("query-auth.enabled", "true") + .build() + .unwrap(), + &format!("{warehouse}/default.db/demo_employees"), + ); + ctx.server.set_auth_response( + "default", + "demo_employees", + AuthTableQueryResponse { + // Row filter: salary >= 90000 (field index 2). + filter: Some(vec![ + r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":2,"name":"salary","type":"BIGINT"}},"function":"GREATER_OR_EQUAL","literals":[90000]}"# + .to_string(), + ]), + // Column masking: name -> UPPER(name) (field index 1). + column_masking: Some(HashMap::from([( + "name".to_string(), + r#"{"name":"UPPER","inputs":[{"index":1,"name":"name","type":"STRING"}]}"# + .to_string(), + )])), + }, + ); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let b = table.new_read_builder(); + // Plan first: scan planning fetches + verifies the per-user grant (mirroring + // Java `CatalogEnvironment.tableQueryAuth()`) and authorizes the shared read + // state; only then may the sync read gate (`new_read`/`to_arrow`) proceed. + let plan = b.new_scan().plan().await.unwrap(); + let enforced: Vec = b + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + dump("ENFORCED (grant: salary>=90000, name->UPPER)", &enforced); + + let mut rows: Vec<(i32, String, i64)> = enforced + .iter() + .flat_map(|batch| { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let names = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let sal = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()) + .map(|r| (ids.value(r), names.value(r).to_string(), sal.value(r))) + .collect::>() + }) + .collect(); + rows.sort_unstable(); + assert_eq!( + rows, + vec![ + (1, "ALICE".to_string(), 120000), + (3, "CHARLIE".to_string(), 95000), + (5, "EVE".to_string(), 99000), + ], + "row filter must drop salary<90000 and masking must uppercase name" + ); +} + +/// Query-auth: scan planning transparently fetches the per-user grant +/// (mirroring Java's `CatalogEnvironment.tableQueryAuth()`); an unrestricted +/// user reads everything, a filtered/masked user gets an enforced read, and +/// paths that cannot enforce stay fail-closed. +#[tokio::test] +async fn test_catalog_get_table_query_auth() { + let ctx = setup_catalog(vec!["default"]).await; + let schema = Schema::builder() + .column("id", DataType::BigInt(BigIntType::new())) + .option("query-auth.enabled", "true") + .build() + .expect("Failed to build schema"); + ctx.server.add_table_with_schema( + "default", + "qa", + schema, + "file:///tmp/test_warehouse/default.db/qa", + ); + let identifier = Identifier::new("default", "qa"); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + + // Reading no splits yields no rows, so it is allowed even before planning; + // real splits carrying no grant fail closed (read_builder unit tests). + let ungranted = table.new_read_builder().new_read().unwrap(); + assert!(ungranted.to_arrow(&[]).is_ok()); + + // The mock /auth endpoint reports unrestricted by default: planning a scan + // authorizes the table and stamps the grant on its splits. + let builder = table.new_read_builder(); + builder + .new_scan() + .plan() + .await + .expect("unrestricted user should be able to plan a query-auth scan"); + + // A parseable row filter (Java Predicate JSON) grants a filtered read: + // planning and building the read succeed; the filter is enforced inside + // `to_arrow`. + ctx.server.set_auth_response( + "default", + "qa", + AuthTableQueryResponse { + filter: Some(vec![ + r#"{"kind":"LEAF","transform":{"name":"FIELD_REF","fieldRef":{"index":0,"name":"id","type":"BIGINT"}},"function":"GREATER_THAN","literals":[5]}"# + .to_string(), + ]), + column_masking: None, + }, + ); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let builder = table.new_read_builder(); + builder + .new_scan() + .plan() + .await + .expect("a parseable row filter should allow planning a (filtered) read"); + // ... but paths that bypass the row filter stay strictly fail-closed. + let err = table + .new_vector_search_builder() + .execute() + .await + .unwrap_err(); + assert!( + err.to_string().contains("query-auth.enabled"), + "search must stay fail-closed for a filtered user, got: {err}" + ); + + // An unparseable filter fails the plan and keeps the table fail-closed. + ctx.server.set_auth_response( + "default", + "qa", + AuthTableQueryResponse { + filter: Some(vec!["{\"kind\":\"CUSTOM\"}".to_string()]), + column_masking: None, + }, + ); + let fresh = ctx.catalog.get_table(&identifier).await.unwrap(); + assert!(fresh.new_read_builder().new_scan().plan().await.is_err()); + + // Parseable column masking grants a (masked) read; a caller predicate on + // the masked column is rejected (it would leak the raw value). + ctx.server.set_auth_response( + "default", + "qa", + AuthTableQueryResponse { + filter: None, + column_masking: Some(HashMap::from([( + "id".to_string(), + "{\"name\":\"NULL\"}".to_string(), + )])), + }, + ); + let table = ctx.catalog.get_table(&identifier).await.unwrap(); + let builder = table.new_read_builder(); + builder.new_scan().plan().await.expect("masked read plans"); + // A caller predicate on a masked column fails closed at plan time (pruning + // on its raw value would leak it); the same guard runs again in `to_arrow`. + let mut filtered = table.new_read_builder(); + filtered.with_filter( + PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Long(1)) + .unwrap(), + ); + let Err(err) = filtered.new_scan().plan().await else { + panic!("a caller predicate on a masked column must fail closed"); + }; + assert!( + err.to_string().contains("masked column"), + "a caller predicate on a masked column must fail closed, got: {err}" + ); + + // Every plan re-authorizes (like Java): revoking down to an unparseable + // grant fails the plan, and a read without a stamped grant stays closed. + ctx.server.set_auth_response( + "default", + "qa", + AuthTableQueryResponse { + filter: Some(vec!["{\"kind\":\"CUSTOM\"}".to_string()]), + column_masking: None, + }, + ); + assert!(table.new_read_builder().new_scan().plan().await.is_err()); +} + #[tokio::test] async fn test_rest_env_get_table_reuses_catalog_environment() { let ctx = setup_catalog(vec!["default"]).await;