[table] Support query-auth row filters and column masking at read time

Author: plusplusjiajia

crates/integrations/datafusion/Cargo.toml

@@ -49,6 +49,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"] }
 flate2 = "1"
 parquet = { workspace = true }
 tar = "0.4"

crates/integrations/datafusion/src/table/mod.rs

@@ -377,7 +377,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()

crates/integrations/datafusion/src/variant_pushdown.rs

@@ -235,6 +235,7 @@ impl ExtensionPlanner for VariantExtractionExtensionPlanner {
                 .collect()
         };
         let filter_exact = !filter_analysis.requires_residual
+            && plan.row_counts_exact()
             && filter_analysis
                 .pushed_predicate
                 .as_ref()

crates/integrations/datafusion/tests/query_auth.rs

@@ -0,0 +1,209 @@
+// 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::{Int32Array, Int64Array, StringArray};
+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::<Int32Array>().unwrap();
+        let names = b.column(1).as_any().downcast_ref::<StringArray>().unwrap();
+        let sal = b.column(2).as_any().downcast_ref::<Int64Array>().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::<Int64Array>()
+        .unwrap()
+        .value(0);
+    assert_eq!(count, 3, "COUNT(*) must not use unfiltered statistics");
+}

crates/paimon/src/api/api_request.rs

@@ -167,6 +167,21 @@ impl AlterTableRequest {
     }
 }
 
+/// Request for auth table query: the projected columns of the query.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct AuthTableQueryRequest {
+    /// Projected column names; `None` means all columns.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub select: Option<Vec<String>>,
+}
+
+impl AuthTableQueryRequest {
+    /// Create a new AuthTableQueryRequest.
+    pub fn new(select: Option<Vec<String>>) -> Self {
+        Self { select }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -193,6 +208,18 @@ mod tests {
         assert!(json.contains("\"updates\""));
     }
 
+    #[test]
+    fn test_auth_table_query_request_serialization() {
+        let req = AuthTableQueryRequest::new(Some(vec!["a".to_string(), "b".to_string()]));
+        assert_eq!(
+            serde_json::to_string(&req).unwrap(),
+            r#"{"select":["a","b"]}"#
+        );
+        // `None` omits the key entirely (matches the server's optional field).
+        let req = AuthTableQueryRequest::new(None);
+        assert_eq!(serde_json::to_string(&req).unwrap(), "{}");
+    }
+
     #[test]
     fn test_rename_table_request_serialization() {
         let source = Identifier::new("db1".to_string(), "table1".to_string());

crates/paimon/src/api/api_response.rs

@@ -427,10 +427,55 @@ pub struct GetTableTokenResponse {
     pub expires_at_millis: Option<i64>,
 }
 
+/// Response for auth table query: the per-user row filter and column masking the
+/// client must enforce at read time for a `query-auth.enabled` table.
+#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct AuthTableQueryResponse {
+    /// JSON-serialized row-filter predicates, ANDed together. Empty/None = no filter.
+    pub filter: Option<Vec<String>>,
+    /// column name -> JSON-serialized masking transform. Empty/None = no masking.
+    pub column_masking: Option<HashMap<String, String>>,
+}
+
+impl AuthTableQueryResponse {
+    /// True when the server imposes no row filter and no column masking.
+    pub fn is_unrestricted(&self) -> bool {
+        self.filter.as_ref().is_none_or(|f| f.is_empty())
+            && self.column_masking.as_ref().is_none_or(|m| m.is_empty())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_auth_table_query_response_deserialization() {
+        // A restricted grant: pin the exact wire field names. A drift in either
+        // (`filter` / `columnMasking`) would deserialize to None and silently
+        // skip authorization, so this test fails closed against that.
+        let resp: AuthTableQueryResponse =
+            serde_json::from_str(r#"{"filter":["p0","p1"],"columnMasking":{"ssn":"m0"}}"#).unwrap();
+        assert_eq!(resp.filter, Some(vec!["p0".to_string(), "p1".to_string()]));
+        assert_eq!(
+            resp.column_masking,
+            Some(HashMap::from([("ssn".to_string(), "m0".to_string())]))
+        );
+        assert!(!resp.is_unrestricted());
+
+        // The real server sends `{}` for an unrestricted grant (both fields are
+        // `@JsonInclude(NON_NULL)` in Java); it must parse to an empty grant.
+        let empty: AuthTableQueryResponse = serde_json::from_str("{}").unwrap();
+        assert_eq!(empty, AuthTableQueryResponse::default());
+        assert!(empty.is_unrestricted());
+
+        // Present-but-empty collections are also unrestricted.
+        let blank: AuthTableQueryResponse =
+            serde_json::from_str(r#"{"filter":[],"columnMasking":{}}"#).unwrap();
+        assert!(blank.is_unrestricted());
+    }
+
     #[test]
     fn test_error_response_serialization() {
         let resp = ErrorResponse::new(

crates/paimon/src/api/mod.rs

@@ -31,16 +31,16 @@ mod api_response;
 
 // Re-export request types
 pub use api_request::{
-    AlterDatabaseRequest, AlterTableRequest, CreateDatabaseRequest, CreateFunctionRequest,
-    CreateTableRequest, CreateViewRequest, RenameTableRequest,
+    AlterDatabaseRequest, AlterTableRequest, AuthTableQueryRequest, CreateDatabaseRequest,
+    CreateFunctionRequest, CreateTableRequest, CreateViewRequest, RenameTableRequest,
 };
 
 // Re-export response types
 pub use api_response::{
-    AuditRESTResponse, ConfigResponse, ErrorResponse, GetDatabaseResponse, GetFunctionResponse,
-    GetTableResponse, GetTableTokenResponse, GetViewResponse, ListDatabasesResponse,
-    ListFunctionsResponse, ListPartitionsResponse, ListTablesResponse, ListViewsResponse,
-    PagedList,
+    AuditRESTResponse, AuthTableQueryResponse, ConfigResponse, ErrorResponse, GetDatabaseResponse,
+    GetFunctionResponse, GetTableResponse, GetTableTokenResponse, GetViewResponse,
+    ListDatabasesResponse, ListFunctionsResponse, ListPartitionsResponse, ListTablesResponse,
+    ListViewsResponse, PagedList,
 };
 
 // Re-export error types

crates/paimon/src/api/resource_paths.rs

@@ -147,6 +147,11 @@ impl ResourcePaths {
         )
     }
 
+    /// Get the auth-table-query endpoint path (`.../tables/{table}/auth`).
+    pub fn auth_table(&self, database_name: &str, table_name: &str) -> String {
+        format!("{}/auth", self.table(database_name, table_name))
+    }
+
     /// Get the table details endpoint path.
     pub fn table_details(&self, database_name: &str) -> String {
         format!(
@@ -242,6 +247,15 @@ mod tests {
         assert!(table_path.contains("my-table"));
     }
 
+    #[test]
+    fn test_auth_table_path_encodes_names() {
+        let paths = ResourcePaths::new("catalog");
+        assert_eq!(
+            paths.auth_table("analytics db", "user events"),
+            "/v1/catalog/databases/analytics+db/tables/user+events/auth"
+        );
+    }
+
     #[test]
     fn test_config_path() {
         assert_eq!(ResourcePaths::config(), "/v1/config");