feat(auth): authorize query-auth reads and carry the grant on the split

crates/integrations/datafusion/src/partition_count_pushdown.rs

@@ -296,7 +296,14 @@ impl TableProvider for PartitionRowCountProvider {
             .clone();
         let table = self.table.clone();
         let table = crate::runtime::await_with_runtime(async move {
-            CoreOptions::new(table.schema().options()).ensure_read_authorized()?;
+            // Rules the manifests cannot apply: stay unpinned, so the scan
+            // carries the grant.
+            if CoreOptions::new(table.schema().options())
+                .ensure_read_authorized()
+                .is_err()
+            {
+                return Ok(Some(table));
+            }
             if table.travel_snapshot().is_some() {
                 return Ok(Some(table));
             }
@@ -328,6 +335,7 @@ impl TableProvider for PartitionRowCountProvider {
             self,
             table,
             provider_as_source(Arc::new(fallback_provider)),
+            provider_as_source(Arc::new(self.fallback_provider.clone())),
             projection.cloned(),
             state,
         )?);
@@ -352,6 +360,7 @@ struct PartitionRowCountStream {
     projection: Option<Vec<usize>>,
     output_schema: SchemaRef,
     source: Arc<dyn TableSource>,
+    unpinned_source: Arc<dyn TableSource>,
     table_name: TableReference,
     filters: Vec<Expr>,
     state: Arc<SessionState>,
@@ -362,6 +371,7 @@ impl PartitionRowCountStream {
         provider: &PartitionRowCountProvider,
         table: Option<Table>,
         source: Arc<dyn TableSource>,
+        unpinned_source: Arc<dyn TableSource>,
         projection: Option<Vec<usize>>,
         state: SessionState,
     ) -> DFResult<Self> {
@@ -374,6 +384,7 @@ impl PartitionRowCountStream {
             projection,
             output_schema,
             source,
+            unpinned_source,
             table_name: provider.table_name.clone(),
             filters: provider.filters.clone(),
             state: Arc::new(state),
@@ -387,13 +398,24 @@ impl PartitionRowCountStream {
         let counts = match self.table.clone() {
             Some(table) => {
                 let predicate = self.predicate.clone();
-                crate::runtime::await_with_runtime(async move {
+                match crate::runtime::await_with_runtime(async move {
                     table
                         .exact_partition_row_counts_with_filter(predicate)
                         .await
                 })
                 .await
-                .map_err(to_datafusion_error)?
+                {
+                    Ok(counts) => counts,
+                    // Refused: only an unpinned scan can authorize.
+                    Err(paimon::Error::Unsupported { .. }) => {
+                        let plan = crate::runtime::await_with_runtime(
+                            self.scan_by_reading(&self.unpinned_source),
+                        )
+                        .await?;
+                        return self.fallback_stream(plan, context);
+                    }
+                    Err(error) => return Err(to_datafusion_error(error)),
+                }
             }
             None => Some(Vec::new()),
         };
@@ -407,15 +429,9 @@ impl PartitionRowCountStream {
                 self.table_name,
                 self.table.as_ref().and_then(Table::travel_snapshot).map(|snapshot| snapshot.id()),
             );
-            let plan = crate::runtime::await_with_runtime(self.scan_by_reading()).await?;
-            if plan.schema() != self.output_schema {
-                return internal_err!(
-                    "partition count fallback schema mismatch: expected {:?}, got {:?}",
-                    self.output_schema,
-                    plan.schema()
-                );
-            }
-            return datafusion::physical_plan::execute_stream(plan, context);
+            let plan =
+                crate::runtime::await_with_runtime(self.scan_by_reading(&self.source)).await?;
+            return self.fallback_stream(plan, context);
         };
 
         // A partition with no surviving rows must not create a GROUP BY key.
@@ -476,16 +492,34 @@ impl PartitionRowCountStream {
     }
 
     /// The same rows, computed the ordinary way: count the original scan per partition.
-    async fn scan_by_reading(&self) -> DFResult<Arc<dyn ExecutionPlan>> {
-        let source_schema = self.source.schema();
+    fn fallback_stream(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        context: Arc<TaskContext>,
+    ) -> DFResult<SendableRecordBatchStream> {
+        if plan.schema() != self.output_schema {
+            return internal_err!(
+                "partition count fallback schema mismatch: expected {:?}, got {:?}",
+                self.output_schema,
+                plan.schema()
+            );
+        }
+        datafusion::physical_plan::execute_stream(plan, context)
+    }
+
+    async fn scan_by_reading(
+        &self,
+        source: &Arc<dyn TableSource>,
+    ) -> DFResult<Arc<dyn ExecutionPlan>> {
+        let source_schema = source.schema();
         let partition_indices = self
             .partition_fields
             .iter()
             .map(|field| source_schema.index_of(field.name()))
             .collect::<Result<Vec<_>, _>>()?;
         let scan = LogicalPlan::TableScan(TableScan::try_new(
             self.table_name.clone(),
-            Arc::clone(&self.source),
+            Arc::clone(source),
             Some(partition_indices),
             self.filters.clone(),
             None,

crates/paimon/src/api/api_response.rs

@@ -568,6 +568,7 @@ impl ListPoliciesResponse {
 
 #[cfg(test)]
 mod tests {
+
     use super::*;
 
     #[test]

crates/paimon/src/spec/core_options.rs

@@ -1063,7 +1063,7 @@ impl<'a> CoreOptions<'a> {
             .and_then(|v| v.parse().ok())
     }
 
-    fn configured_time_travel_selectors(&self) -> Vec<&'static str> {
+    pub(crate) fn configured_time_travel_selectors(&self) -> Vec<&'static str> {
         let mut selectors = Vec::with_capacity(6);
         if self.options.contains_key(SCAN_TIMESTAMP_MILLIS_OPTION) {
             selectors.push(SCAN_TIMESTAMP_MILLIS_OPTION);

crates/paimon/src/spec/schema.rs

@@ -638,22 +638,27 @@ impl TableSchema {
     }
 }
 
-/// Reject column names reserved for system use, mirroring Java `SpecialFields`:
-/// the five `SYSTEM_FIELD_NAMES` and the `_KEY_` key-field prefix.
+/// Whether `name` is one Paimon reserves for a system column. Java
+/// `SpecialFields.SYSTEM_FIELD_NAMES` plus the `_KEY_` key-field prefix.
+pub(crate) fn is_reserved_system_field_name(name: &str) -> bool {
+    name.starts_with(KEY_FIELD_PREFIX) || SYSTEM_FIELD_NAMES.contains(&name)
+}
+
+// Java SpecialFields.SYSTEM_FIELD_NAMES.
+const SYSTEM_FIELD_NAMES: [&str; 5] = [
+    SEQUENCE_NUMBER_FIELD_NAME,
+    VALUE_KIND_FIELD_NAME,
+    "_LEVEL",
+    ROW_KIND_FIELD_NAME,
+    ROW_ID_FIELD_NAME,
+];
+const KEY_FIELD_PREFIX: &str = "_KEY_";
+
+/// Reject column names reserved for system use, mirroring Java `SpecialFields`.
 ///
 /// A user column colliding with a system field is otherwise excluded from the
 /// physical read and silently filled with the system value.
 fn validate_no_reserved_field_names(fields: &[DataField]) -> crate::Result<()> {
-    // Java SpecialFields.SYSTEM_FIELD_NAMES.
-    const SYSTEM_FIELD_NAMES: [&str; 5] = [
-        SEQUENCE_NUMBER_FIELD_NAME,
-        VALUE_KIND_FIELD_NAME,
-        "_LEVEL",
-        ROW_KIND_FIELD_NAME,
-        ROW_ID_FIELD_NAME,
-    ];
-    const KEY_FIELD_PREFIX: &str = "_KEY_";
-
     for field in fields {
         let name = field.name();
         if name.starts_with(KEY_FIELD_PREFIX) || SYSTEM_FIELD_NAMES.contains(&name) {

crates/paimon/src/table/audit_log_table/read.rs

@@ -44,7 +44,9 @@ pub struct AuditLogRead<'a> {
 
 impl<'a> AuditLogRead<'a> {
     pub fn new(read: TableRead<'a>) -> crate::Result<Self> {
-        read.ensure_query_auth_allowed()?;
+        // Query-auth is decided per split in `to_arrow`; only the type is known here.
+        crate::spec::CoreOptions::new(read.table().schema().options())
+            .ensure_type_paimon_served(&read.table().identifier().full_name())?;
         match read.0 {
             TableReadKind::Paimon(read) => Ok(Self { read }),
             TableReadKind::Format(_) => Err(crate::Error::Unsupported {
@@ -55,6 +57,9 @@ impl<'a> AuditLogRead<'a> {
 
     /// Reads splits planned by an audit scan, retaining winning retract rows.
     pub fn to_arrow(&self, data_splits: &[DataSplit]) -> crate::Result<ArrowRecordBatchStream> {
+        // The primary-key path below builds its readers directly, so decide here.
+        self.read
+            .ensure_authorized_by_splits(&self.read.table.schema.core_options(), data_splits)?;
         let output_read_type = self.read.read_type.clone();
         if output_read_type
             .iter()

crates/paimon/src/table/format_table_read.rs

@@ -118,7 +118,15 @@ impl<'a> FormatTableRead<'a> {
         data_splits: &[DataSplit],
     ) -> crate::Result<ArrowRecordBatchStream> {
         let core_options = self.table.schema().core_options();
-        core_options.ensure_read_authorized()?;
+        core_options.ensure_type_paimon_served(&self.table.identifier().full_name())?;
+        // The marker carries the plan's decision.
+        if core_options.query_auth_enabled()
+            || data_splits.iter().any(|split| split.query_auth_required())
+        {
+            return Err(super::query_auth::unsupported(
+                "a format table cannot apply a row filter or column masking",
+            ));
+        }
         // Mapping the conjunct onto the data fields drops it, so the read would
         // silently ignore the filter. Guard on the read path, not the builder:
         // `TableRead` is public and can be constructed and filtered directly.

crates/paimon/src/table/mod.rs

@@ -107,6 +107,7 @@ mod postpone_fixed_bucket_router;
 mod postpone_fixed_bucket_write;
 mod postpone_fixed_bucket_write_builder;
 mod prepared_files;
+mod query_auth;
 mod read_builder;
 pub mod referenced_files;
 pub(crate) mod rest_env;
@@ -221,6 +222,9 @@ pub struct Table {
     schema_manager: SchemaManager,
     branch: String,
     branch_reference: bool,
+    /// Minted only by [`RESTEnv::build_table`]; a [`Table::new`] handle cannot
+    /// replay a grant.
+    query_auth_session: Option<u64>,
     rest_env: Option<RESTEnv>,
     /// True when this table copy was switched to a historical schema by
     /// [`Table::copy_with_time_travel`]. Such a copy is read-only.
@@ -250,6 +254,7 @@ impl Table {
             schema_manager,
             branch,
             branch_reference: false,
+            query_auth_session: None,
             rest_env,
             time_traveled: false,
             travel_snapshot: None,
@@ -290,6 +295,7 @@ impl Table {
             schema_manager,
             branch,
             branch_reference,
+            query_auth_session: None,
             rest_env: None,
             time_traveled: false,
             travel_snapshot: None,
@@ -374,6 +380,75 @@ impl Table {
         }
     }
 
+    /// Whether this handle reads a schema other than the one the server rules
+    /// on: a time-travel option, a travelled or branch view, or a `$branch_x` /
+    /// `$files` name that reads the base table's files.
+    pub(crate) fn reads_another_schema(&self) -> Result<bool> {
+        // Presence only; planning validates the selector after adapting
+        // `scan.version`.
+        let travels = !CoreOptions::new(self.schema.options())
+            .configured_time_travel_selectors()
+            .is_empty();
+        let decorated = self.identifier.branch_name()?.is_some()
+            || self.identifier.system_table_name()?.is_some();
+        Ok(travels || self.time_traveled || self.branch_reference || decorated)
+    }
+
+    /// Whether this user may read this table; `None` when it is not
+    /// `query-auth.enabled`. `query_auth` is the loaded option, as in Java.
+    pub(crate) async fn authorize_read(
+        &self,
+        query_auth: bool,
+    ) -> Result<Option<std::sync::Arc<query_auth::QueryAuthGrant>>> {
+        let local = CoreOptions::new(self.schema.options());
+        let Some(rest_env) = &self.rest_env else {
+            // Only a REST catalog can authorize.
+            return if local.query_auth_enabled() {
+                Err(query_auth::unsupported(
+                    "it requires a REST catalog to authorize the query",
+                ))
+            } else {
+                Ok(None)
+            };
+        };
+
+        // No freshness assertion yet — an ordinary table must not inherit one.
+        if !query_auth {
+            return Ok(None);
+        }
+        if self.reads_another_schema()? {
+            return Err(query_auth::unsupported(
+                "a time-travelled or branch read authorizes against the table's current schema, \
+                 which is not the one it reads",
+            ));
+        }
+
+        // Before any RPC: only the catalog mints a session, so an assembled
+        // handle stops here.
+        let session = self.query_auth_session.ok_or_else(|| {
+            query_auth::unsupported("this table handle was assembled rather than loaded")
+        })?;
+
+        // Naming a system column here would fail the server's column check.
+        let response = rest_env
+            .table_query_auth(self.schema.id(), self.schema.fields(), None)
+            .await?;
+        Ok(Some(std::sync::Arc::new(query_auth::QueryAuthGrant::new(
+            response, session,
+        ))))
+    }
+
+    /// Handed out once per catalog-loaded table; wraps only after 2^64 loads.
+    pub(crate) fn with_query_auth_session(mut self) -> Self {
+        static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
+        self.query_auth_session = Some(NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed));
+        self
+    }
+
+    pub(crate) fn query_auth_session(&self) -> Option<u64> {
+        self.query_auth_session
+    }
+
     /// Get the REST environment, if this table was loaded from a REST catalog.
     pub fn rest_env(&self) -> Option<&RESTEnv> {
         self.rest_env.as_ref()
@@ -483,6 +558,7 @@ impl Table {
             schema_manager: self.schema_manager.clone(),
             branch: self.branch.clone(),
             branch_reference: self.branch_reference,
+            query_auth_session: self.query_auth_session,
             rest_env: self.rest_env.clone(),
             time_traveled: self.time_traveled,
             travel_snapshot: if selector_changed {
@@ -517,6 +593,8 @@ impl Table {
             branch_reference: self.branch_reference || branch != DEFAULT_MAIN_BRANCH,
             time_traveled: false,
             travel_snapshot: None,
+            // Not the schema the catalog loaded, so not a handle it authorizes.
+            query_auth_session: None,
             ..self.clone()
         })
     }
@@ -724,6 +802,7 @@ impl Table {
             schema_manager,
             branch,
             branch_reference: true,
+            query_auth_session: self.query_auth_session,
             rest_env: self.rest_env.clone(),
             time_traveled: false,
             travel_snapshot: None,
@@ -758,6 +837,32 @@ pub(crate) fn find_field_id_by_name(fields: &[DataField], name: &str) -> Option<
     fields.iter().find(|f| f.name() == name).map(|f| f.id())
 }
 
+/// A `query-auth.enabled` table wired to its own REST session.
+#[cfg(test)]
+pub(crate) async fn rest_query_auth_table() -> Table {
+    use crate::api::rest_api::RESTApi;
+    use crate::common::{CatalogOptions, Options};
+
+    let mut options = Options::default();
+    options.set(CatalogOptions::URI, "http://127.0.0.1:1");
+    options.set("token.provider", "bear");
+    options.set("token", "test_token");
+    let api = std::sync::Arc::new(RESTApi::new(options.clone(), false).await.unwrap());
+    let table = query_auth_table();
+    Table {
+        rest_env: Some(RESTEnv::new(
+            table.identifier.clone(),
+            "uuid-1".to_string(),
+            api,
+            options,
+            false,
+            None,
+        )),
+        ..table
+    }
+    .with_query_auth_session()
+}
+
 /// A minimal table with `query-auth.enabled = true`, for the fail-closed read guard.
 #[cfg(test)]
 pub(crate) fn query_auth_table() -> Table {