[python] Support query auth (row filter & column masking) for REST catalog#8136
[python] Support query auth (row filter & column masking) for REST catalog#8136MgjLLL wants to merge 1 commit into
Conversation
|
I found a few correctness issues in the query-auth paths introduced here:
|
|
Fixes for issues raised by @JingsongLi, plus one additional issue found during analysis. Fix 1: Ray read path bypasses auth (
|
@JingsongLi All 3 issues fixed (+ 1 additional parallel path bypass found during analysis). See updated PR description. PTAL. |
| if not self.filter and not self.column_masking: | ||
| return plan | ||
| auth_splits = [QueryAuthSplit(split, self) for split in plan.splits()] | ||
| return Plan(auth_splits) |
There was a problem hiding this comment.
Java TableScan.Plan does not carry a snapshot id, but Python Plan does and the update / row-id update paths use it as check_from_snapshot. Wrapping the plan here drops plan.snapshot_id, so a query-auth table planned from a non-empty snapshot becomes snapshot_id=None; table_update then emits commit messages with -1, which disables the row-id conflict checks (and related global-index update checks). Please preserve the original plan metadata, e.g. Plan(auth_splits, snapshot_id=plan.snapshot_id).
|
|
||
| return reader | ||
|
|
||
| def _create_split_read_with_read_type(self, split, read_type): |
There was a problem hiding this comment.
This auth-specific construction bypasses the normal PK read path above. In _create_split_read, PK tables inject missing sequence.field columns into the inner read type and then project them back out, matching the Java withReadType + outer projection behavior. Here, if query auth is enabled and the user projects id,val from a PK table with sequence.field=ts, MergeFileSplitRead is built without ts; that can either fail with sequence.field ... not found or merge by file sequence instead of the configured user sequence. Please reuse the existing _create_split_read widening/project-back logic for effective_read_type, or factor it so the auth path cannot drift from the normal PK path.
| elif function == "LIKE": | ||
| raw = literals[0] | ||
| escaped = re.escape(raw) | ||
| pattern = escaped.replace("%", ".*").replace("_", ".") |
There was a problem hiding this comment.
This does not match the JVM LIKE semantics. Java treats backslash as the default escape character before expanding % / _, so a policy predicate like LIKE admin\\_% matches admin_foo and not adminXfoo. Escaping the whole string first and then replacing every % / _ makes escaped wildcards behave as wildcards (or requires a literal backslash), so Python can allow/deny different rows from the Java client for the same auth filter. Please port the Java Like.sqlToRegexLike behavior, including invalid escape handling.
|
@JingsongLi All 3 issues fixed (+ several related correctness issues found during follow-up self-review). New commit Fixes for the three review comments1.
|
MgjLLL
force-pushed
the
python-query-auth
branch
2 times, most recently
from
bf32ff2 to
15ee1c1
Compare
JingsongLi
left a comment
JingsongLi
left a comment
There was a problem hiding this comment.
I found a few correctness issues in the Python query-auth implementation.
| elif function == "IN": | ||
| return pc.is_in(value_array, pa.array(converted, type=value_array.type)) | ||
| elif function == "NOT_IN": | ||
| return pc.invert(pc.is_in(value_array, pa.array(converted, type=value_array.type))) |
There was a problem hiding this comment.
Java Paimon's NotIn returns false when the input field is null, but pc.is_in(null, ...) is false and inverting it makes null rows pass. For a row-filter auth rule such as dept NOT_IN ('blocked'), Python would expose rows where dept is null. Please combine this with pc.is_valid(value_array) and also preserve Java's null-literal behavior, where any null literal makes NOT_IN false.
| self._parsed_rules = {col: json.loads(tj) for col, tj in masking_rules.items()} | ||
| read_field_names = {f.name for f in read_fields} | ||
| for col_name, transform in self._parsed_rules.items(): | ||
| for ref_name in _collect_all_field_refs_from_transform(transform): |
There was a problem hiding this comment.
This validates references for every masking rule before checking whether the masked target column is actually projected. If REST returns a rule like secret = FIELD_REF(email) and the user reads only id, the Python reader raises because email is absent even though secret is not returned. Java skips masking rules whose target column is absent from the output row type before remapping inputs, so this should filter to projected target columns first.
| return pa.array([True] * batch_len, type=pa.bool_()) | ||
| elif function == "FALSE": | ||
| return pa.array([False] * batch_len, type=pa.bool_()) | ||
| raise ValueError(f"Unknown leaf function: {function}") |
There was a problem hiding this comment.
IS_NAN is a valid Java Paimon predicate (PredicateBuilder.isNaN / IsNaN.NAME) and can be serialized in REST auth filters. With the current switch it fails every Python read with Unknown leaf function: IS_NAN; please add a branch using pc.is_nan for float/double arrays.
JingsongLi
left a comment
JingsongLi
left a comment
There was a problem hiding this comment.
I left a few inline comments on query-auth correctness and serialization edge cases.
| elif function == "IS_NOT_NULL": | ||
| return pc.is_valid(value_array) | ||
| elif function == "IN": | ||
| return pc.is_in(value_array, pa.array(converted, type=value_array.type)) |
There was a problem hiding this comment.
Java's In.test returns false when the field value is null and skips null literals. pc.is_in treats NULL IN (..., NULL) as true, so an auth row filter like dept IN ('eng', NULL) would leak rows where dept is null. Please drop null literals and wrap the result with pc.if_else(pc.is_valid(value_array), ..., False); if no non-null literals remain, the result should be all false.
There was a problem hiding this comment.
Fixed. Now drop null literals from the value set and wrap with pc.if_else(pc.is_valid(value_array), in_mask, False). If all literals are null, returns all-false. Also added null literal guards for EQUAL/NOT_EQUAL/LT/LE/GT/GE to match Java's LeafBinaryFunction semantics.
| if not options.query_auth_enabled or self.catalog_loader is None: | ||
| return None | ||
|
|
||
| def auth(select): |
There was a problem hiding this comment.
This nested auth function is stored on FileStoreTable as _query_auth_fn, which makes auth-enabled tables/read tasks unpickleable (AttributeError: Can't get local object ...). Ray/Daft and the existing serialization paths pickle tables or read tasks, so this breaks those paths once query-auth.enabled is true. Please use a top-level callable/dataclass or reconstruct the catalog auth function lazily instead of storing a local closure.
There was a problem hiding this comment.
Fixed. Replaced the nested closure with a top-level _TableQueryAuthFn class that is pickleable. Verified with pickle roundtrip test.
| return pa.nulls(len(original_batch), type=pa.string()) | ||
| sep = self._resolve_input(inputs[0], original_batch) | ||
| values = [self._resolve_input(inp, original_batch) for inp in inputs[1:]] | ||
| return pc.binary_join_element_wise(*values, sep, null_handling='skip') |
There was a problem hiding this comment.
binary_join_element_wise(..., null_handling='skip') does not match Java BinaryString.concatWs when the separator is non-null and all value inputs for a row are null. Java returns an empty string for that row, but Arrow returns a shorter array, so set_column can fail with ArrowInvalid: Added column's length must match record batch's length. Please normalize this case so the result always has len(original_batch) rows; the predicate-side CONCAT_WS path has the same issue.
There was a problem hiding this comment.
Fixed. PyArrow's binary_join_element_wise with null_handling='skip' returns an empty array (length 0) when all values are null — confirmed by testing. Now pad the result to batch length and use pc.if_else(sep_null, None, pc.if_else(pc.is_valid(result), result, "")) to match Java's BinaryString.concatWs behavior: all-null values → empty string, null separator → null. Same fix applied in the predicate transform path.
| from pypaimon.schema.data_types import DataField | ||
|
|
||
|
|
||
| class TableNoPermissionException(Exception): |
There was a problem hiding this comment.
This defines a second TableNoPermissionException that is unrelated to pypaimon.catalog.catalog_exception.TableNoPermissionException. RESTCatalog.auth_table_query raises this new type on 403, so callers catching the catalog exception will miss query-auth denials. Please reuse the existing catalog exception or make this subclass it.
There was a problem hiding this comment.
Fixed. Removed the duplicate class from table_query_auth.py. Updated catalog_exception.TableNoPermissionException to accept an optional cause parameter and handle both Identifier objects and plain strings. All callers now use the unified exception.
MgjLLL
force-pushed
the
python-query-auth
branch
from
034f3b3 to
040a6bf
Compare
| return pa.nulls(len(batch), type=pa.string()) | ||
| result = pc.binary_join_element_wise(*values, sep, null_handling='skip') | ||
| if len(result) < len(batch): | ||
| padded = result.to_pylist() + [None] * (len(batch) - len(result)) |
There was a problem hiding this comment.
Padding the compacted Arrow result at the end misaligns rows. pc.binary_join_element_wise(..., null_handling="skip") drops rows whose value inputs are all null, so for a=[None, "x", None], b=[None, "y", "z"], sep="-" it returns ["x-y", "z"]; this code pads that to ["x-y", "z", ""], while Java BinaryString.concatWs should evaluate the rows as ["", "x-y", "z"]. In a row-filter predicate that can allow or deny the wrong rows. Please make the CONCAT_WS transform row-preserving, and add a mixed all-null/non-null regression test.
| values = [self._resolve_input(inp, original_batch) for inp in inputs[1:]] | ||
| result = pc.binary_join_element_wise(*values, sep, null_handling='skip') | ||
| if len(result) < len(original_batch): | ||
| padded = result.to_pylist() + [None] * (len(original_batch) - len(result)) |
There was a problem hiding this comment.
This has the same row-alignment problem as the predicate CONCAT_WS path. Arrow returns a compacted array when some rows have all-null value inputs, so appending padding at the end shifts the later rows masked values. For example a=[None, "x", None], b=[None, "y", "z"] with sep="-" becomes ["x-y", "z", ""] here, but Java BinaryString.concatWs would produce ["", "x-y", "z"]. That masks the wrong rows; please preserve original row positions and cover the mixed all-null/non-null case in tests.
| for col_idx, masked_array in masked_columns.items(): | ||
| original_field = original_batch.schema.field(col_idx) | ||
| if masked_array.type != original_field.type: | ||
| masked_array = pc.cast(masked_array, original_field.type) |
There was a problem hiding this comment.
Column masking should not force the transformed value back to the original column type. Java TableQueryAuthResult.transform writes transform.transform(row) directly; RESTCatalogTest.testColumnMaskingApplyOnRead masks an INT column with CastTransform(..., STRING) and then reads "100" with getString(2). Here the same {"name":"CAST","type":"STRING"} mask on an int column is cast back to int32, so the mask is effectively undone or can fail for non-castable masking expressions. Please replace the column using the masked array type instead, and update the schema-preservation test accordingly.
MgjLLL
force-pushed
the
python-query-auth
branch
from
040a6bf to
aad7740
Compare
| def extract_row_filter(self) -> Optional[Callable[[pa.RecordBatch], pa.Array]]: | ||
| if not self.filter: | ||
| return None | ||
| filters = [parse_predicate_to_batch_filter(json_str) for json_str in self.filter] |
There was a problem hiding this comment.
Java TableQueryAuthResult.extractPredicate() skips empty filter strings before deserializing them, but this path tries to parse every entry. If the REST server returns filter=[""] (or a mix of empty and valid filters), Python will fail the read with JSONDecodeError in both extract_row_filter() and get_extra_fields_for_filter(), while the JVM client treats the empty entry as no-op. Please filter out empty/blank JSON strings before wrapping the plan or before parsing, and add a regression test for empty filter entries.
There was a problem hiding this comment.
Fixed. Now filter empty/blank entries at the source — TableQueryAuthResult constructor strips falsy entries from the filter list ([f for f in filter if f]) and empty keys/values from column_masking dict. AuthMaskingReader also guards if not tj: continue before json.loads. Python truthiness check (if f) matches Java StringUtils.isEmpty semantics: skips null and "" but passes through whitespace strings. Added regression tests for mixed empty/valid filters and empty masking values.
|
|
||
| def new_read_builder(self) -> 'ReadBuilder': | ||
| return ReadBuilder(self) | ||
| def new_read_builder(self, *, _skip_auth=False) -> 'ReadBuilder': |
There was a problem hiding this comment.
The Python write paths (e.g. table_update.py, table_upsert_by_key.py,
file_store_write.py) internally call new_read_builder() to read existing
data for operations like UPDATE, UPSERT, and MERGE INTO. This is a pre-existing
design that predates this PR.
Once query auth is injected into new_read_builder(), those internal reads
would also be subject to row filtering and column masking — returning a subset
of rows or masked column values to the write engine, which would corrupt the
write results.
_skip_auth=True bypasses auth for these internal, system-initiated reads,
matching Java's design: Java's write paths operate at a lower abstraction level
(e.g. LocalTableQuery, SortedGlobalIndexBuilder) that never passes through
AbstractDataTableScan's auth injection point in the first place.
The leading underscore signals that this is an internal API not intended for
external callers.
JingsongLi
left a comment
JingsongLi
left a comment
There was a problem hiding this comment.
I think we need to take a closer look here. This feature cannot affect the Public API.
- Why modify daft?
- Why modify read_builder?
This is not just questions. This is the tendency of design, and I need you to provide clear solutions instead of why you are currently changing them.
Java has not affected the Public API.
|
AI cannot solve these problems, please provide feedback in the form of a real person. |
…talog
When query-auth.enabled is true, the Python client calls POST
/v1/.../databases/{db}/tables/{tb}/auth before producing a Plan and
applies row-level filters and column masking transforms on the read
path, matching the Java client behavior.
Auth is injected via post-construction on TableScan/StreamingTableScan
by ReadBuilder/StreamReadBuilder — no Public API signatures change.
Write paths disable auth before plan() to preserve data integrity.
New files:
- catalog/table_query_auth.py: TableQueryAuthResult
- common/predicate_json_parser.py: JSON predicate to PyArrow filters
- read/query_auth_split.py: QueryAuthSplit wrapper
- read/reader/auth_masking_reader.py: filter, masking, projection readers
| auth_result = self._auth_query() | ||
| plan = self.file_scanner.scan() | ||
| if auth_result is not None: | ||
| plan = auth_result.convert_plan(plan) |
There was a problem hiding this comment.
This auth wrapping is only applied by plan(), but scan_with_stats() below still returns the raw scanner plan. ReadBuilder.explain() uses scan_with_stats(), and Daft explain derives routing from those split details; the actual Daft scan path sees QueryAuthSplit and falls back to pypaimon, while explain will still report native Parquet for the same auth-enabled scan. Please apply the same _auth_query().convert_plan(...) logic in scan_with_stats() while preserving the ScanStats, and add an explain regression for an auth-enabled table.
2 participants
Purpose
Adds query-auth support to the Python client so it honors the row-level filter and column masking rules returned by a REST catalog, matching the existing JVM client behavior.
When the new option
query-auth.enabledis set totrue, before producing aPlanthe client callsPOST /v1/.../databases/{db}/tables/{tb}/authwith the projected fields, receives{ filter, columnMasking }, and applies them on the read path:RESTApi.auth_table_queryissues the call (new request/response modelsAuthTableQueryRequest/AuthTableQueryResponse, new path inResourcePaths.auth_table).TableQueryAuth/TableQueryAuthResult(catalog/table_query_auth.py) wrap the result and convert each split to aQueryAuthSplit.predicate_json_parser(common/predicate_json_parser.py) parses Paimon predicate JSON into a PyArrow compute filter (EQ/NEQ/LT/LTEQ/GT/GTEQ/IS_NULL/IS_NOT_NULL/IN/NOT_IN/STARTS_WITH/ENDS_WITH/CONTAINS/AND/OR/NOT).AuthFilterReader/AuthMaskingReader/ColumnProjectReader(read/reader/auth_masking_reader.py) implement row filtering, column masking transforms (NULL,FIELD_REF,CAST,UPPER,LOWER,CONCAT,CONCAT_WS) and final projection back to the user's requested columns.read_builder/stream_read_builder/table_read/table_scan/file_store_table/catalog_environment/rest_catalogare wired to invoke the auth call and pull extra fields required only by the auth filter.Behavior is gated by the new
CoreOptions.QUERY_AUTH_ENABLED(query-auth.enabled, defaultfalse), so existing users see no change.Tests
Three new test files (994+ lines, all passing locally under
pytest):paimon-python/pypaimon/tests/predicate_json_parser_test.py— covers each predicate kind, nested AND/OR/NOT, type coercion, null handling, andextract_referenced_fields.paimon-python/pypaimon/tests/auth_masking_reader_test.py— covers each masking transform, missing-field validation, and projection back to the user-requested columns.paimon-python/pypaimon/tests/table_query_auth_test.py— end-to-end coverage: REST catalog callsauth_table_query, the result is plumbed into the plan, splits becomeQueryAuthSplit, and reads return filtered + masked rows.Local check:
API and Format
query-auth.enabled(boolean, defaultfalse).POST /v1/{prefix}/databases/{db}/tables/{tb}/auth. Request{ "select": [...] }, response{ "filter": [<predicate-json>...], "columnMasking": { <col>: <transform-json>, ... } }. The contract follows the existing Java client; no server-side change is required for catalogs that already implement query auth.AuthTableQueryRequest,AuthTableQueryResponse,TableQueryAuth,TableQueryAuthResult,QueryAuthSplit,AuthFilterReader,AuthMaskingReader,ColumnProjectReader) are additive and live under existing modules.Documentation
The new option
query-auth.enabledshould be reflected in the Python configuration reference. Happy to add the docs entry in this PR or in a follow-up — please advise.This closes #8135