diff --git a/paimon-python/pypaimon/common/predicate_json_parser.py b/paimon-python/pypaimon/common/predicate_json_parser.py index a6474908ef36..df14ce209731 100644 --- a/paimon-python/pypaimon/common/predicate_json_parser.py +++ b/paimon-python/pypaimon/common/predicate_json_parser.py @@ -172,14 +172,14 @@ def _apply_predicate_transform(transform: dict, batch: pa.RecordBatch, elif name in _DATE_EXTRACT: return _date_extract(name, transform["fieldRef"], batch) elif name == "LENGTH": - _check_string_input("LENGTH input", transform["inputs"][0], batch) + _check_string_input("LENGTH input", transform["inputs"][0], batch, stored=True) source = _resolve_transform_input(transform["inputs"][0], batch) return pa.array( [None if v is None else len(v) for v in source.to_pylist()], type=pa.int32()) elif name == "BIT_LENGTH": - _check_string_input("BIT_LENGTH input", transform["inputs"][0], batch) + _check_string_input("BIT_LENGTH input", transform["inputs"][0], batch, stored=True) source = _resolve_transform_input(transform["inputs"][0], batch) return pa.array( [None if v is None else len(v.encode("utf-8")) * 8 @@ -320,6 +320,24 @@ def _check_string_input(slot: str, value, batch: pa.RecordBatch, stored: bool = raise ValueError(f"{slot} field must be a string: {column.type}") +_INTEGER_TYPE = re.compile( + r"\s*(?:TINYINT|SMALLINT|INT|INTEGER|BIGINT)(?=\s|\Z)" + r"(?:\s*(?:NOT\s+)?NULL)?\s*\Z", re.IGNORECASE) + + +def _check_integer_input(slot, value, batch): + if value is None or (isinstance(value, int) and not isinstance(value, bool)): + return + if not isinstance(value, dict): + raise ValueError(f"{slot} must be an integer or a field: {value!r}") + stored = value.get("type") + if not isinstance(stored, str) or _INTEGER_TYPE.match(stored) is None: + raise ValueError(f"{slot} field must be an integer: {stored!r}") + column = _field_column(value, batch) + if not pa.types.is_integer(column.type): + raise ValueError(f"{slot} field must be an integer: {column.type}") + + class _Positions: """A position slot resolved lazily, as Java reads one only when a row reaches it.""" @@ -432,8 +450,9 @@ def _pad(inputs, direction, batch: pa.RecordBatch) -> pa.Array: raise ValueError(f"PAD takes 3 inputs, got {inputs!r}") if direction not in ("LEFT", "RIGHT"): raise ValueError(f"PAD direction must be LEFT or RIGHT: {direction!r}") - _check_string_input("PAD input", inputs[0], batch) - _check_string_input("PAD string", inputs[2], batch) + _check_string_input("PAD input", inputs[0], batch, stored=True) + _check_string_input("PAD string", inputs[2], batch, stored=True) + _check_integer_input("PAD length", inputs[1], batch) source = _resolve_transform_input(inputs[0], batch) pad = _resolve_transform_input(inputs[2], batch) lengths = _Positions(inputs[1], batch) @@ -444,10 +463,11 @@ def _pad(inputs, direction, batch: pa.RecordBatch) -> pa.Array: value = src_list[i] pad_str = pad_list[i] raw_len = lengths.value(i) - if value is None or raw_len is None or pad_str is None: + length = _int_position(raw_len) if raw_len is not None else None + if value is None or length is None or pad_str is None: result.append(None) continue - result.append(_pad_one(value, _int_position(raw_len), pad_str, direction)) + result.append(_pad_one(value, length, pad_str, direction)) return pa.array(result, type=pa.string()) @@ -464,9 +484,9 @@ def _pad_one(source, length, pad, direction): def _translate(inputs, batch: pa.RecordBatch) -> pa.Array: if not isinstance(inputs, list) or len(inputs) != 3: raise ValueError(f"TRANSLATE takes 3 inputs, got {inputs!r}") - _check_string_input("TRANSLATE source", inputs[0], batch) - _check_string_input("TRANSLATE matching", inputs[1], batch) - _check_string_input("TRANSLATE replacement", inputs[2], batch) + _check_string_input("TRANSLATE source", inputs[0], batch, stored=True) + _check_string_input("TRANSLATE matching", inputs[1], batch, stored=True) + _check_string_input("TRANSLATE replacement", inputs[2], batch, stored=True) source = _resolve_transform_input(inputs[0], batch) matching = _resolve_transform_input(inputs[1], batch) replacement = _resolve_transform_input(inputs[2], batch) @@ -499,8 +519,11 @@ def _translate_one(source, matching, replacement): def _overlay(inputs, batch: pa.RecordBatch) -> pa.Array: if not isinstance(inputs, list) or len(inputs) not in (3, 4): raise ValueError(f"OVERLAY takes 3 or 4 inputs, got {inputs!r}") - _check_string_input("OVERLAY input", inputs[0], batch) - _check_string_input("OVERLAY replacement", inputs[1], batch) + _check_string_input("OVERLAY input", inputs[0], batch, stored=True) + _check_string_input("OVERLAY replacement", inputs[1], batch, stored=True) + _check_integer_input("OVERLAY position", inputs[2], batch) + if len(inputs) == 4: + _check_integer_input("OVERLAY length", inputs[3], batch) source = _resolve_transform_input(inputs[0], batch) replacement = _resolve_transform_input(inputs[1], batch) positions = _Positions(inputs[2], batch) @@ -513,15 +536,13 @@ def _overlay(inputs, batch: pa.RecordBatch) -> pa.Array: value = src_list[i] repl = repl_list[i] raw_pos = positions.value(i) - # SQL null propagation: any null argument yields a null result - if value is None or repl is None or raw_pos is None: - result.append(None) - continue - if has_length and lengths.value(i) is None: + # Java evaluates both integer arguments before propagating string NULLs. + pos = _int_position(raw_pos) if raw_pos is not None else None + raw_length = lengths.value(i) if has_length else None + length = _int_position(raw_length) if raw_length is not None else None + if value is None or repl is None or pos is None or (has_length and length is None): result.append(None) continue - pos = _int_position(raw_pos) - length = _int_position(lengths.value(i)) if has_length else None replaced = len(repl) if (length is None or length < 0) else length # pos - 1 and pos + replaced are 32-bit int expressions in Java's # OverlayTransform; wrap them so a boundary position overflows exactly diff --git a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py index 22dfc43ab672..e2630f72f44d 100644 --- a/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/filter_record_batch_reader.py @@ -92,7 +92,13 @@ def _filter_batch(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: ) def _filter_batch_by_row(self, batch: pa.RecordBatch) -> Optional[pa.RecordBatch]: - df = polars.from_arrow(pa.Table.from_batches([batch])) + # Predicate indices address logical fields. A record-reader adapter + # may prepend _row_kind, which is metadata rather than an operand. + columns = ([field.name for field in self.schema_fields] + if self.schema_fields is not None else + [name for name in batch.schema.names if name != "_row_kind"]) + logical_batch = batch.select(columns) + df = polars.from_arrow(pa.Table.from_batches([logical_batch])) iterator = InternalRowWrapperIterator( self._iter_df_rows(df), df.width, diff --git a/paimon-python/pypaimon/read/split_serializer.py b/paimon-python/pypaimon/read/split_serializer.py index 0dbf0484b261..e28197a67ab5 100644 --- a/paimon-python/pypaimon/read/split_serializer.py +++ b/paimon-python/pypaimon/read/split_serializer.py @@ -18,7 +18,8 @@ """Read and write the cross-language ``SplitSerializer`` v1 binary format. This is the stable bridge between Java, PyPaimon and ``pypaimon_rust``. It -supports DataSplit and IndexedSplit (including scores on the Python side) and +supports DataSplit, IndexedSplit (including scores on the Python side), +QueryAuthSplit with the captured rules, and the current DataSplit v9/DataFileMeta v21 layout. The reader remains compatible with DataSplit v8. """ @@ -31,6 +32,7 @@ from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.read.split import DataSplit, Split +from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.schema.data_types import AtomicType, DataField from pypaimon.table.row.binary_row import BinaryRow from pypaimon.table.row.generic_row import ( @@ -44,6 +46,7 @@ _SPLIT_SER_VERSION = 1 _TYPE_DATA_SPLIT = 1 _TYPE_INDEXED_SPLIT = 3 +_TYPE_QUERY_AUTH_SPLIT = 5 _INDEXED_SPLIT_MAGIC = -938472394838495695 _INDEXED_SPLIT_VERSION = 1 _SPLIT_MAGIC = -2394839472490812314 @@ -290,7 +293,11 @@ def serialize_split_v1(split: Split, include_scores: bool = True) -> bytes: writer = _Writer() writer.i64(_SPLIT_SER_MAGIC) writer.i32(_SPLIT_SER_VERSION) - if isinstance(split, IndexedSplit): + if isinstance(split, QueryAuthSplit): + writer.i32(_TYPE_QUERY_AUTH_SPLIT) + writer.take(serialize_split_v1(split.split, include_scores)) + _write_auth_result(writer, split.auth_result) + elif isinstance(split, IndexedSplit): data_split = split.data_split() if not isinstance(data_split, DataSplit): raise TypeError( @@ -326,6 +333,32 @@ def serialize_split_v1(split: Split, include_scores: bool = True) -> bytes: return writer.finish() +def _write_auth_result(writer, result): + writer.u8(result is not None) + if result is None: + return + + def string(value): + if value is None: + writer.i32(-1) + else: + encoded = value.encode('utf-8') + writer.i32(len(encoded)) + writer.take(encoded) + + writer.u8(result.filter is not None) + if result.filter is not None: + writer.i32(len(result.filter)) + for rule in result.filter: + string(rule) + writer.u8(result.column_masking is not None) + if result.column_masking is not None: + writer.i32(len(result.column_masking)) + for column in sorted(result.column_masking, key=lambda name: name.encode('utf-16-be')): + string(column) + string(result.column_masking[column]) + + def _serialized_row(row) -> bytes: if row is None: row = GenericRow([], []) @@ -527,6 +560,14 @@ def deserialize_split_v1(data: bytes, partition_fields: List[DataField], merge-on-read; None for append tables. """ r = _Reader(data) + split = _read_split_frame(r, partition_fields, key_fields) + r.finish() + return split + + +def _read_split_frame(r, partition_fields, key_fields, depth=0): + if depth >= 32: + raise ValueError('QueryAuthSplit nesting exceeds 32') magic = r.i64() if magic != _SPLIT_SER_MAGIC: raise ValueError("bad SplitSerializer magic %d" % magic) @@ -535,7 +576,11 @@ def deserialize_split_v1(data: bytes, partition_fields: List[DataField], raise ValueError( "unsupported SplitSerializer version %d (expected %d)" % (version, _SPLIT_SER_VERSION)) type_id = r.i32() - if type_id == _TYPE_DATA_SPLIT: + if type_id == _TYPE_QUERY_AUTH_SPLIT: + split = QueryAuthSplit( + _read_split_frame(r, partition_fields, key_fields, depth + 1), + _read_auth_result(r)) + elif type_id == _TYPE_DATA_SPLIT: split = _read_datasplit_body(r, partition_fields, key_fields) elif type_id == _TYPE_INDEXED_SPLIT: imagic = r.i64() @@ -562,10 +607,38 @@ def deserialize_split_v1(data: bytes, partition_fields: List[DataField], split = IndexedSplit(data_split, row_ranges, scores) else: raise ValueError("unsupported split type id %d" % type_id) - r.finish() return split +def _read_auth_result(reader): + from pypaimon.catalog.table_query_auth import TableQueryAuthResult + + if reader.u8() == 0: + return None + + def count(multiplier): + size = reader.i32() + if size < 0 or size > (len(reader.d) - reader.p) // (4 * multiplier): + raise ValueError('invalid QueryAuthSplit rule count') + return size + + def string(): + size = reader.i32() + if size < 0: + raise ValueError('NULL authorization rule in QueryAuthSplit') + return reader.take(size).decode('utf-8') + + filters = [string() for _ in range(count(1))] if reader.u8() else None + masks = None + if reader.u8(): + masks = {} + for _ in range(count(2)): + # Python 3.6/3.7 dict comprehensions evaluate values before keys. + column = string() + masks[column] = string() + return TableQueryAuthResult(filters, masks) + + def _read_scores(r: '_Reader') -> Optional[List[float]]: if r.u8() == 0: return None diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index b0b103be4a80..a7188cd9c98a 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -29,13 +29,15 @@ from pypaimon.common.predicate import Predicate from pypaimon.common.predicate_json_parser import extract_referenced_fields from pypaimon.data.map_shared_shredding import map_selected_keys_field -from pypaimon.read.push_down_utils import predicate_field_names +from pypaimon.read.push_down_utils import ( + exclude_predicate_with_fields, predicate_field_names, rewrite_predicate_indices) from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.reader.auth_masking_reader import ( AuthFilterReader, AuthMaskingReader, ColumnProjectReader, RecordReaderToBatchAdapter, BatchToRecordReaderAdapter) from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.read.reader.limited_record_reader import LimitedRecordBatchReader +from pypaimon.read.reader.filter_record_batch_reader import FilterRecordBatchReader from pypaimon.read.split import Split from pypaimon.read.variant_read_type import has_variant_extractions from pypaimon.read.read_type import (adapter_output_paths, extract_array, output_schema, reader_adapter) @@ -48,6 +50,7 @@ from pypaimon.utils.arrow_utils import zero_column_batch ROW_KIND_COLUMN = "_row_kind" +_DEFAULT_PREDICATE = object() logger = logging.getLogger(__name__) _RECORD_BATCH_READER_FROM_STREAM = getattr( pyarrow.ipc.RecordBatchReader, "from_stream", None) @@ -500,11 +503,13 @@ def _try_native_batches( if not self._native_blob_view_supported(): return self._native_fallback( "the table uses an unsupported BLOB view") - # Query authorization has additional filtering, masking and projection - # semantics which are already implemented by the Python reader. - if any(isinstance(split, QueryAuthSplit) for split in splits): + # Only an opaque split from Rust's authorized plan retains the live + # grant. Python-planned/serialized rules remain on the Python reader; + # reconstructing a Rust split cannot mint authorization from bytes. + if any(isinstance(split, QueryAuthSplit) and getattr(split, '_native_split', None) is None + for split in splits): return self._native_fallback( - "query authorization requires the Python reader") + "query authorization requires the original native plan") try: from pypaimon.read.native_plan import ( _prepare_native_read, _raise_if_native_fork_safety_error, @@ -1457,7 +1462,8 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, push_down_limit: bool = True, post_merge_filter=None, eager_blob_fields=None, - post_filter_after_inline: bool = False) -> SplitRead: + post_filter_after_inline: bool = False, + predicate=_DEFAULT_PREDICATE) -> SplitRead: sr = self._build_split_read( split, read_type, @@ -1466,6 +1472,7 @@ def _create_split_read(self, split: Split, blob_parallelism: int = 1, post_merge_filter, eager_blob_fields, post_filter_after_inline, + predicate, ) sr._blob_parallelism = blob_parallelism sr._parquet_row_group_cache = self._parquet_row_group_cache @@ -1476,7 +1483,9 @@ def _build_split_read(self, split: Split, read_type=None, push_down_limit: bool = True, post_merge_filter=None, eager_blob_fields=None, - post_filter_after_inline: bool = False) -> SplitRead: + post_filter_after_inline: bool = False, + predicate=_DEFAULT_PREDICATE) -> SplitRead: + predicate = self.predicate if predicate is _DEFAULT_PREDICATE else predicate effective_limit = ( self.limit if limit is None else limit ) if push_down_limit else None @@ -1529,7 +1538,7 @@ def _build_split_read(self, split: Split, read_type=None, outer_extract_name_paths = self._output_extract_name_paths() return MergeFileSplitRead( table=self.table, - predicate=self.predicate, + predicate=predicate, read_type=inner_read_type, split=split, row_tracking_enabled=False, @@ -1556,7 +1565,7 @@ def _build_split_read(self, split: Split, read_type=None, outer_extract_name_paths = self._output_extract_name_paths() return DataEvolutionSplitRead( table=self.table, - predicate=self.predicate, + predicate=predicate, read_type=scan_read_type, split=split, row_tracking_enabled=True, @@ -1588,7 +1597,7 @@ def _build_split_read(self, split: Split, read_type=None, outer_extract_name_paths = self._output_extract_name_paths() return RawFileSplitRead( table=self.table, - predicate=self.predicate, + predicate=predicate, read_type=inner_read_type, split=split, row_tracking_enabled=self.table.options.row_tracking_enabled(), @@ -1801,12 +1810,17 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, auth_result.validate_read_type( latest_fields, read_fields, self.nested_name_paths, table_fields) - extra_fields = auth_result.get_extra_fields(read_fields, table_fields) - effective_read_type = read_fields - if extra_fields: - effective_read_type = read_fields + extra_fields + # Query filters can reference masked columns absent from the output. + # Include them and their mask dependencies before the final projection. + filter_fields = self._predicate_extra_fields + needed_fields = read_fields + filter_fields + auth_extra_fields = auth_result.get_extra_fields(needed_fields, table_fields) + extra_fields = filter_fields + auth_extra_fields + effective_read_type = read_fields + extra_fields filter_fn = auth_result.extract_row_filter() + raw_predicate = exclude_predicate_with_fields( + self.predicate, set(auth_result.column_masking or {})) effective_limit = self.limit if limit is None else limit auth_fields = ( self._auth_filter_field_names(auth_result, effective_read_type) @@ -1828,7 +1842,8 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, blob_parallelism=blob_parallelism, read_type=effective_read_type, limit=limit, - push_down_limit=filter_fn is None or embed_filter, + push_down_limit=(filter_fn is None or embed_filter) and self.predicate is None, + predicate=raw_predicate, post_merge_filter=filter_fn if embed_filter else None, eager_blob_fields=auth_fields if embed_filter else None, post_filter_after_inline=post_filter_after_inline, @@ -1852,12 +1867,17 @@ def __authed_reader(self, split, auth_result, blob_parallelism=1, if filter_fn and not embed_filter: reader = AuthFilterReader(reader, filter_fn) - if effective_limit is not None: - reader = LimitedRecordBatchReader(reader, effective_limit) if auth_result.column_masking: reader = AuthMaskingReader(reader, auth_result.column_masking, effective_read_type) + if self.predicate is not None: + predicate = rewrite_predicate_indices(self.predicate, effective_read_type) + reader = FilterRecordBatchReader( + reader, predicate, schema_fields=effective_read_type) + if effective_limit is not None: + reader = LimitedRecordBatchReader(reader, effective_limit) + if extra_fields: original_columns = [f.name for f in read_fields] reader = ColumnProjectReader(reader, original_columns) diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index eb9c06d7ef5e..771b2872958c 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -27,7 +27,7 @@ from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.manifest.manifest_list_manager import ManifestListManager from pypaimon.read.plan import Plan -from pypaimon.read.query_auth_split import resolve_auth_result, wrap_plan_with_auth +from pypaimon.read.query_auth_split import QueryAuthSplit, resolve_auth_result, wrap_plan_with_auth from pypaimon.read.scan_stats import ScanStats from pypaimon.read.scanner.file_scanner import FileScanner @@ -65,10 +65,18 @@ def __init__( self.file_scanner = self._create_file_scanner() def plan(self) -> Plan: + native_auth = (self.table.options.query_auth_enabled + and self._native_requested() and self._native_plan_supported()) + if native_auth: + # Rust asks the REST server and carries its complete response in + # QueryAuthSplit. Both readers then enforce that captured policy. + native = self._try_native_plan() + if native is not None: + return native auth_result = self.__auth_query() # The native planner bypasses the auth-aware file scanner. Resolve auth # before selecting a planning backend. - if (auth_result is None and self._native_requested() + if (not native_auth and auth_result is None and self._native_requested() and self._native_plan_supported()): native = self._try_native_plan() if native is not None: @@ -135,7 +143,16 @@ def _native_plan_supported_impl(self) -> bool: if not database_name or database_name == UNKNOWN_DATABASE: return False if self.table.options.query_auth_enabled: - return False + # The binding currently carries explicit partition selectors as + # ordinary filters. Keep them on Python, where masked scan-only + # selectors can be rejected as Java requires. + if self.partition_predicate is not None: + return False + # The filesystem path has no catalog-issued authorization session. + loader = getattr(getattr(self.table, 'catalog_environment', None), 'catalog_loader', None) + from pypaimon.read.native_plan import _catalog_metastore + if loader is None or _catalog_metastore(loader) != 'rest': + return False # Rust rejects schemas whose primary keys are all partition keys. if getattr(self.table, 'is_primary_key_table', False) \ and not self.table.trimmed_primary_keys: @@ -234,7 +251,9 @@ def _try_native_plan(self) -> Optional[Plan]: # physical order need not match the PK merge comparator. return None partition_predicate = self.file_scanner.partition_key_predicate - if partition_predicate is not None: + if partition_predicate is not None and not any(isinstance(split, QueryAuthSplit) for split in splits): + # Authorized user predicates observe masked partition values; + # Rust has already applied only the safe raw partition tests. splits = [s for s in splits if getattr(s, 'partition', None) is None or partition_predicate.test(s.partition)] @@ -274,7 +293,8 @@ def plan_for_write(self) -> Plan: self.file_scanner.skip_level0 = skip_level0 def __auth_query(self): - return authorize(self.table, self._query_auth_fn, self._read_type) + return authorize(self.table, self._query_auth_fn, self._read_type, + self.predicate, self.partition_predicate) def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]: """Run :meth:`plan` while recording manifest / pruning counters. @@ -283,8 +303,14 @@ def scan_with_stats(self) -> Tuple[Plan, Optional[ScanStats]]: keeps going through :meth:`plan`. Native plans return stats=None; explain reports their split metadata without pruning counters. """ + native_auth = (self.table.options.query_auth_enabled + and self._native_requested() and self._native_plan_supported()) + if native_auth: + native = self._try_native_plan() + if native is not None: + return native, None auth_result = self.__auth_query() - if (auth_result is None and self._native_requested() + if (not native_auth and auth_result is None and self._native_requested() and self._native_plan_supported()): native = self._try_native_plan() if native is not None: @@ -581,14 +607,43 @@ def validate_auth_rules(table, auth_result, latest_fields=None): return latest_fields -def authorize(table, query_auth_fn, read_type): - auth_result = resolve_auth_result(query_auth_fn, read_type) +def authorize(table, query_auth_fn, read_type, predicate=None, partition_predicate=None): + if query_auth_fn is None: + return None + requested_type = read_type + if read_type: + from pypaimon.read.push_down_utils import predicate_field_names + selected = {field.name for field in read_type} + names = set(selected) + for query_filter in (predicate, partition_predicate): + if query_filter is not None: + names.update(predicate_field_names(query_filter)) + requested_type = list(read_type) + [ + field for field in table.fields if field.name in names and field.name not in selected] + auth_result = resolve_auth_result(query_auth_fn, requested_type) if auth_result is not None: validate_auth_rules(table, auth_result) + if partition_predicate is not None and auth_result.column_masking: + from pypaimon.read.push_down_utils import predicate_field_names + if predicate_field_names(partition_predicate) & set(auth_result.column_masking): + raise ValueError("Explicit partition filters cannot reference masked columns") return auth_result def prune_scanner_by_auth(table, scanner, auth_result): + # Java excludes masked columns from all raw-data pruning. The complete + # query predicate is executed after masks have produced visible values. + if auth_result.column_masking: + from pypaimon.read.push_down_utils import exclude_predicate_with_fields + masked = set(auth_result.column_masking) + query_predicate = getattr(scanner, 'predicate', None) + for name in ('predicate', 'predicate_for_stats', 'primary_key_predicate', 'partition_key_predicate'): + setattr(scanner, name, exclude_predicate_with_fields(getattr(scanner, name, None), masked)) + init_selector = getattr(scanner, '_init_bucket_selector', None) + if init_selector is not None: + scanner._bucket_selector = init_selector() + if query_predicate is not None: + scanner.auth_has_non_partition_filter = True if not auth_result.filter: return partition_preds, has_non_partition = __split_auth_filter(table, auth_result) diff --git a/paimon-python/pypaimon/tests/native_query_auth_test.py b/paimon-python/pypaimon/tests/native_query_auth_test.py new file mode 100644 index 000000000000..74503ccda7bf --- /dev/null +++ b/paimon-python/pypaimon/tests/native_query_auth_test.py @@ -0,0 +1,302 @@ +# 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. + +"""Existing REST authorization rules through Python and Native readers.""" + +import json +import pickle +from contextlib import ExitStack +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon import Schema +from pypaimon.api.api_response import AuthTableQueryResponse +from pypaimon.common.predicate_builder import PredicateBuilder +from pypaimon.read.query_auth_split import QueryAuthSplit +from pypaimon.read.split_serializer import deserialize_split_v1, serialize_split_v1 +from pypaimon.read.native_plan import native_runtime_available +from pypaimon.schema.schema_change import SchemaChange +from pypaimon.tests import native_plan_rest_test + +rest_catalog = native_plan_rest_test.rest_catalog +pytestmark = [pytest.mark.native_plan, pytest.mark.skipif( + not native_runtime_available(), reason='Rust reader required')] + + +def _field(name, type_='STRING'): + return {'name': name, 'type': type_, 'index': 0} + + +@pytest.fixture(params=[ + ('append', {}), + ('primary-key', {'bucket': '1'}), + ('data-evolution', {'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}), +], ids=lambda mode: mode[0]) +def guarded_table(rest_catalog, request): + catalog, server = rest_catalog + schema = pa.schema([('id', pa.int32()), ('value', pa.string()), ('source', pa.string()), + ('position', pa.int64()), ('length', pa.int32()), ('n', pa.int32())]) + mode, options = request.param + catalog.create_table('default.t', Schema.from_pyarrow_schema( + schema, primary_keys=['id'] if mode == 'primary-key' else None, options=options), False) + table = catalog.get_table('default.t') + writer = table.new_batch_write_builder().new_write() + commit = table.new_batch_write_builder().new_commit() + try: + writer.write_arrow(pa.table({ + 'id': [1, 2, 3, 4], 'value': ['secret'] * 4, + 'source': ['η”²πŸ™‚δΉ™', 'abcabc', '', None], + 'position': [2, -3, 0, 1], 'length': [2, 2, 0, None], 'n': [0] * 4, + }, schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + catalog.alter_table(table.identifier, [SchemaChange.set_option('query-auth.enabled', 'true')]) + return catalog.get_table('default.t'), server + + +def _read(table, native, projection=None, predicate=None, limit=None): + table = table.copy({'read.native.enabled': str(native).lower(), + 'scan.native-plan.enabled': str(native).lower()}) + builder = table.new_read_builder() + if projection is not None: + builder.with_projection(projection) + if predicate is not None: + builder.with_filter(predicate) + if limit is not None: + builder.with_limit(limit) + scan = builder.new_scan() + with ExitStack() as stack: + if native: + stack.enter_context(patch.object(scan.file_scanner, 'scan', + side_effect=AssertionError('Native plan fell back'))) + stack.enter_context(patch('pypaimon.read.table_read.TableRead._create_split_read', + side_effect=AssertionError('Native read fell back'))) + plan = scan.plan() + return builder.new_read().to_arrow(plan.splits()) + + +@pytest.mark.parametrize('rule,expected', [ + ({'name': 'SUBSTRING', 'inputs': [_field('source'), 2, 2]}, ['πŸ™‚δΉ™', 'bc', '', None]), + ({'name': 'SUBSTRING', + 'inputs': [_field('source'), _field('position', 'BIGINT'), _field('length', 'INT')]}, + ['πŸ™‚δΉ™', 'ab', '', None]), + ({'name': 'TRIM', 'trimFlag': 'BOTH', 'inputs': [_field('source'), 'η”²δΉ™ac']}, + ['πŸ™‚', 'bcab', '', None]), + ({'name': 'PAD', 'direction': 'LEFT', 'inputs': [_field('source'), 5, 'πŸ™‚x']}, + ['πŸ™‚xη”²πŸ™‚δΉ™', 'abcab', 'πŸ™‚xπŸ™‚xπŸ™‚', None]), + ({'name': 'TRANSLATE', 'inputs': [_field('source'), 'η”²πŸ™‚δΉ™aba', 'XY']}, ['XY', 'cc', '', None]), + ({'name': 'OVERLAY', 'inputs': [_field('source'), 'x', 2, 1]}, ['η”²xδΉ™', 'axcabc', 'x', None]), +]) +def test_existing_string_masks_read_natively(guarded_table, rule, expected): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'value': json.dumps(rule)}) + for native in [False, True]: + result = _read(table, native, ['id', 'value']).sort_by([('id', 'ascending')]) + assert result.column('value').to_pylist() == expected + + +@pytest.mark.parametrize('name,expected', [('LENGTH', [3, 6, 0, None]), ('BIT_LENGTH', [80, 48, 0, None])]) +def test_integer_masks_use_hidden_string_dependencies(guarded_table, name, expected): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'n': json.dumps({'name': name, 'inputs': [_field('source')]})}) + for native in [False, True]: + result = _read(table, native, ['id', 'n']).sort_by([('id', 'ascending')]) + assert result.column('n').type == pa.int32() + assert result.column('n').to_pylist() == expected + + +def _filter(field, function, literals, type_='INT'): + return json.dumps({'kind': 'LEAF', 'transform': {'name': 'FIELD_REF', 'fieldRef': _field(field, type_)}, + 'function': function, 'literals': literals}) + + +def test_auth_filter_precedes_masks_and_user_filter_precedes_limit(guarded_table): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + filter=[_filter('id', 'GREATER_THAN', [1]), _filter('source', 'IS_NOT_NULL', [], 'STRING')], + column_masking={'value': json.dumps({'name': 'SUBSTRING', 'inputs': [_field('source'), 1, 2]})}) + predicate = PredicateBuilder(table.fields).equal('value', 'ab') + for native in [False, True]: + result = _read(table, native, ['id'], predicate, limit=1) + assert result.to_pylist() == [{'id': 2}] + + +@pytest.mark.parametrize('mode', ['disable-native', 'setup-failure', 'wire-roundtrip', 'pickle']) +def test_native_plan_retains_captured_policy_on_python_read_fallback(guarded_table, mode): + table, server = guarded_table + response = AuthTableQueryResponse( + filter=[_filter('id', 'EQUAL', [2])], + column_masking={'value': json.dumps( + {'name': 'PAD', 'direction': 'RIGHT', 'inputs': [_field('source'), 3, 'x']})}) + server.table_query_auth_store[table.identifier.get_full_name()] = response + table = table.copy({'scan.native-plan.enabled': 'true', 'read.native.enabled': 'true'}) + builder = table.new_read_builder().with_projection(['id', 'value']) + scan = builder.new_scan() + with patch.object(scan.file_scanner, 'scan', side_effect=AssertionError('Native plan fell back')): + splits = scan.plan().splits() + assert all(isinstance(split, QueryAuthSplit) for split in splits) + assert splits + # A policy re-fetch would be observably wrong. Fallback must execute the + # response captured by this plan, even if the server now grants all rows. + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse() + if mode == 'disable-native': + builder = table.copy({'read.native.enabled': 'false'}).new_read_builder().with_projection(['id', 'value']) + elif mode == 'wire-roundtrip': + splits = [deserialize_split_v1(serialize_split_v1(split), []) for split in splits] + elif mode == 'pickle': + splits = pickle.loads(pickle.dumps(splits)) + with ExitStack() as stack: + if mode == 'setup-failure': + stack.enter_context(patch('pypaimon.read.native_plan.native_read', + side_effect=RuntimeError('setup failed'))) + result = builder.new_read().to_arrow(splits) + assert result.to_pylist() == [{'id': 2, 'value': 'abc'}] + + +def test_masked_partition_predicate_does_not_prune_raw_partition(rest_catalog): + catalog, server = rest_catalog + schema = pa.schema([('id', pa.int32()), ('p', pa.string())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema(schema, partition_keys=['p']), False) + table = catalog.get_table('default.t') + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.table({'id': [1, 2], 'p': ['raw', 'another']}, schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + catalog.alter_table(table.identifier, [SchemaChange.set_option('query-auth.enabled', 'true')]) + table = catalog.get_table('default.t') + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'p': json.dumps({'name': 'CONCAT', 'inputs': ['*']})}) + predicate = PredicateBuilder(table.fields).equal('p', '*') + for native in [False, True]: + result = _read(table, native, ['id', 'p'], predicate).sort_by([('id', 'ascending')]) + assert result.to_pylist() == [{'id': 1, 'p': '*'}, {'id': 2, 'p': '*'}] + + +@pytest.mark.parametrize('rule,expected', [ + ({'name': 'SUBSTRING', 'inputs': ['abc', 'Ω’', 1]}, 'b'), + ({'name': 'SUBSTRING', 'inputs': [None, 'bad']}, None), + ({'name': 'SUBSTRING', 'inputs': ['abc', '99999999999999999999', None]}, None), + ({'name': 'SUBSTRING', 'inputs': [None, 18446744073709551617]}, None), +]) +def test_substring_literal_parsing_obeys_java_null_checks(guarded_table, rule, expected): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'value': json.dumps(rule)}) + for native in [False, True]: + assert _read(table, native, ['value']).column('value').to_pylist() == [expected] * 4 + + +@pytest.mark.parametrize('native', [False, True]) +def test_auth_request_includes_hidden_query_fields(guarded_table, native): + table, server = guarded_table + selections = [] + + def authorize(select): + selections.append(select) + return AuthTableQueryResponse(column_masking={ + 'value': json.dumps({'name': 'CONCAT', 'inputs': ['visible']}) + } if select is None or 'value' in select else None) + + server.table_query_auth_store[table.identifier.get_full_name()] = authorize + predicate = PredicateBuilder(table.fields).equal('value', 'visible') + result = _read(table, native, ['id'], predicate).sort_by([('id', 'ascending')]) + assert result.column('id').to_pylist() == [1, 2, 3, 4] + assert selections and all(set(select) == {'id', 'value'} for select in selections) + + +def test_setup_fallback_applies_masked_hidden_predicate_before_limit(guarded_table): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'value': json.dumps({'name': 'SUBSTRING', 'inputs': [_field('source'), 1, 2]})}) + table = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + builder = (table.new_read_builder().with_projection(['id']).with_limit(1) + .with_filter(PredicateBuilder(table.fields).equal('value', 'ab'))) + scan = builder.new_scan() + with patch.object(scan.file_scanner, 'scan', side_effect=AssertionError('Native plan fell back')): + splits = scan.plan().splits() + with patch('pypaimon.read.native_plan.native_read', side_effect=RuntimeError('setup failed')): + assert builder.new_read().to_arrow(splits).to_pylist() == [{'id': 2}] + + +@pytest.mark.parametrize('native', [False, True]) +def test_explicit_partition_selector_rejects_a_masked_operand(rest_catalog, native): + catalog, server = rest_catalog + schema = pa.schema([('id', pa.int32()), ('p', pa.string())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema( + schema, partition_keys=['p'], options={'query-auth.enabled': 'true'}), False) + table = catalog.get_table('default.t').copy({'scan.native-plan.enabled': str(native).lower()}) + selections = [] + + def authorize(select): + selections.append(select) + return AuthTableQueryResponse(column_masking={'p': json.dumps({'name': 'CONCAT', 'inputs': ['*']})}) + + server.table_query_auth_store[table.identifier.get_full_name()] = authorize + builder = (table.new_read_builder().with_projection(['id']) + .with_partition_filter(PredicateBuilder(table.fields).equal('p', 'raw'))) + with pytest.raises(ValueError, match='partition filters cannot reference masked columns'): + builder.new_scan().plan() + assert selections and set(selections[0]) == {'id', 'p'} + + +def test_pk_row_kind_does_not_shift_a_hidden_string_filter(rest_catalog): + from pypaimon.read.table_read import TableRead + catalog, server = rest_catalog + schema = pa.schema([('id', pa.int32()), ('secret', pa.string())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema( + schema, primary_keys=['id'], options={'bucket': '1'}), False) + table = catalog.get_table('default.t') + for values in [['old-a', 'old-b'], ['new-a', 'new-b']]: + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.table({'id': [1, 2], 'secret': values}, schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + catalog.alter_table(table.identifier, [SchemaChange.set_option('query-auth.enabled', 'true')]) + table = catalog.get_table('default.t').copy({'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false'}) + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'secret': json.dumps({'name': 'CONCAT', 'inputs': ['visible']})}) + predicate = PredicateBuilder(table.fields).contains('secret', 'vis') + builder = table.new_read_builder().with_projection(['id']).with_filter(predicate) + plan = builder.new_scan().plan() + assert any(not split.raw_convertible for split in plan.splits()) + reader = TableRead(table, predicate, builder.read_type(), include_row_kind=True) + rows = reader.to_arrow(plan.splits()).sort_by([('id', 'ascending')]).to_pylist() + assert rows == [{'_row_kind': '+I', 'id': 1}, {'_row_kind': '+I', 'id': 2}] + + +def test_explain_uses_the_authorized_native_plan(guarded_table): + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + column_masking={'value': json.dumps({'name': 'TRIM', 'trimFlag': 'BOTH', 'inputs': [_field('source')]})}) + table = table.copy({'scan.native-plan.enabled': 'true'}) + with patch('pypaimon.read.scanner.file_scanner.FileScanner.scan_with_stats', + side_effect=AssertionError('Native explain fell back')): + assert table.new_read_builder().with_projection(['value']).explain().native_planned diff --git a/paimon-python/pypaimon/tests/predicate_json_parser_test.py b/paimon-python/pypaimon/tests/predicate_json_parser_test.py index fba06c35fcff..7008a40e4d07 100644 --- a/paimon-python/pypaimon/tests/predicate_json_parser_test.py +++ b/paimon-python/pypaimon/tests/predicate_json_parser_test.py @@ -943,6 +943,31 @@ def test_pad_left_right_and_truncate(self): ["hello"]), ["hel"]) + def test_integer_overflow_is_eager_even_with_null_string_inputs(self): + for rule in [ + {"name": "PAD", "direction": "LEFT", "inputs": [None, 2147483648, "x"]}, + {"name": "OVERLAY", "inputs": [None, "x", 2147483648, None]}, + {"name": "OVERLAY", "inputs": [None, "x", None, 2147483648]}, + ]: + with self.subTest(rule=rule), self.assertRaises(ValueError): + self._apply(rule, [None]) + + def test_field_types_are_checked_before_rebinding(self): + batch = pa.RecordBatch.from_pydict({"s": ["abc"], "n": [2]}) + for name, inputs in [ + ("PAD", [self._field(), {"name": "n", "type": "STRING"}, "x"]), + ("OVERLAY", [self._field(), "x", {"name": "n", "type": "STRING"}]), + ("TRANSLATE", [{"name": "s", "type": "INT"}, "a", "b"]), + ("LENGTH", [{"name": "s", "type": "INT"}]), + ("BIT_LENGTH", [{"name": "s", "type": "INT"}]), + ]: + rule = {"name": name, "inputs": inputs, "direction": "LEFT"} + with self.subTest(rule=rule), self.assertRaises(ValueError): + _apply_predicate_transform(rule, batch) + for value in ["2", True, 1.5]: + with self.subTest(value=value), self.assertRaises(ValueError): + self._apply({"name": "PAD", "inputs": [None, value, "x"], "direction": "LEFT"}, [None]) + def test_length_wired_through_a_leaf_filter(self): pred = json.dumps({ "kind": "LEAF", diff --git a/paimon-python/pypaimon/tests/query_auth_validation_test.py b/paimon-python/pypaimon/tests/query_auth_validation_test.py index b61eb78152d2..9b78dbd19dee 100644 --- a/paimon-python/pypaimon/tests/query_auth_validation_test.py +++ b/paimon-python/pypaimon/tests/query_auth_validation_test.py @@ -109,6 +109,8 @@ def __init__(self, fields, latest_fields=None, latest_error=None): def _make_scan(auth_result, read_type, table): scan = TableScan.__new__(TableScan) + scan.predicate = None + scan.partition_predicate = None scan._query_auth_fn = lambda select: auth_result scan._read_type = read_type scan.table = table diff --git a/paimon-python/pypaimon/tests/resources/split_v1_query_auth.bin b/paimon-python/pypaimon/tests/resources/split_v1_query_auth.bin new file mode 100644 index 000000000000..3978a4398cb1 Binary files /dev/null and b/paimon-python/pypaimon/tests/resources/split_v1_query_auth.bin differ diff --git a/paimon-python/pypaimon/tests/rest/rest_server.py b/paimon-python/pypaimon/tests/rest/rest_server.py index b7769eb22f33..f8ce233d6995 100755 --- a/paimon-python/pypaimon/tests/rest/rest_server.py +++ b/paimon-python/pypaimon/tests/rest/rest_server.py @@ -22,6 +22,7 @@ import uuid from http.server import BaseHTTPRequestHandler, HTTPServer from pathlib import Path +import json from typing import Any, Dict, List, Optional, Tuple, Union, TYPE_CHECKING from urllib.parse import urlparse @@ -224,6 +225,7 @@ def __init__(self, data_path: str, auth_provider, config: ConfigResponse, wareho self.no_permission_databases: List[str] = [] self.no_permission_tables: List[str] = [] self.table_token_store: Dict[str, "RESTToken"] = {} + self.table_query_auth_store = {} # Initialize mock catalog (simplified) self.data_path = data_path @@ -558,6 +560,8 @@ def _handle_table_resource(self, method: str, path_parts: List[str], return self._table_commit_handle(method, data, lookup_identifier, branch_part) elif operation == "token": return self._table_token_handle(method, lookup_identifier) + elif operation == "auth": + return self._table_query_auth_handle(method, data, lookup_identifier) elif operation == "rollback": return self._table_rollback_handle(method, data, lookup_identifier) elif operation == "snapshot": @@ -1096,6 +1100,17 @@ def _table_handle(self, method: str, data: str, identifier: Identifier, return self._mock_response(ErrorResponse(None, None, "Method Not Allowed", 405), 405) + def _table_query_auth_handle(self, method, data, identifier): + if method != "POST": + return self._mock_response(ErrorResponse(None, None, "Method Not Allowed", 405), 405) + if identifier.get_full_name() not in self.table_metadata_store: + raise TableNotExistException(identifier) + from pypaimon.api.api_response import AuthTableQueryResponse + response = self.table_query_auth_store.get(identifier.get_full_name(), AuthTableQueryResponse()) + if callable(response): + response = response(json.loads(data).get('select')) + return self._mock_response(response, 200) + def _table_token_handle(self, method: str, identifier: Identifier) -> Tuple[str, int]: if method != "GET": return self._mock_response(ErrorResponse(None, None, "Method Not Allowed", 405), 405) diff --git a/paimon-python/pypaimon/tests/split_serializer_test.py b/paimon-python/pypaimon/tests/split_serializer_test.py index 21533e857e58..ff8cd6f1433b 100644 --- a/paimon-python/pypaimon/tests/split_serializer_test.py +++ b/paimon-python/pypaimon/tests/split_serializer_test.py @@ -18,8 +18,11 @@ import base64 import struct import unittest +from pathlib import Path from pypaimon.globalindex.indexed_split import IndexedSplit +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.split_serializer import ( _decode_modified_utf8, _decode_str_array, _encode_modified_utf8, deserialize_split_v1, serialize_split_v1) @@ -106,6 +109,71 @@ def _partition_fields(self): return [DataField(0, 'y', AtomicType('INT')), DataField(1, 'm', AtomicType('INT'))] + def test_query_auth_java_golden_roundtrips_exactly(self): + frame = (Path(__file__).parent / 'resources' / 'split_v1_query_auth.bin').read_bytes() + split = deserialize_split_v1(frame, self._partition_fields(), + key_fields=[DataField(0, 'k', AtomicType('BIGINT'))]) + self.assertIsInstance(split, QueryAuthSplit) + self.assertEqual(split.auth_result.filter, ['rowη”²πŸ™‚\x00']) + self.assertEqual(split.auth_result.column_masking, {'𐀀': 'supplementary', '\ue000': 'bmp'}) + self.assertEqual(serialize_split_v1(split), frame) + + def test_query_auth_matches_java_utf8_and_utf16_map_order(self): + # Captured from Java QueryAuthSplit: int-length UTF-8 strings, including + # a literal NUL, and TreeMap ordering by UTF-16 rather than code point. + tail = bytes.fromhex( + "0101000000010000000b726f77e794b2f09f9982000100000002" + "00000004f09080800000000d737570706c656d656e74617279" + "00000003ee808000000003626d70") + inner = deserialize_split_v1(_GOLDEN_DATA_SPLIT_V1, self._partition_fields()) + auth = TableQueryAuthResult(["rowη”²πŸ™‚\x00"], {"\ue000": "bmp", "𐀀": "supplementary"}) + frame = serialize_split_v1(QueryAuthSplit(inner, auth)) + expected = struct.pack('>qii', 0x53504C49545F5631, 1, 5) + serialize_split_v1(inner) + tail + self.assertEqual(frame, expected) + copy = deserialize_split_v1(frame, self._partition_fields()) + self.assertIsInstance(copy, QueryAuthSplit) + self.assertEqual(copy.auth_result.filter, auth.filter) + self.assertEqual(copy.auth_result.column_masking, auth.column_masking) + self.assertEqual(copy.row_count, inner.row_count) + self.assertIsNone(copy.merged_row_count()) + self.assertEqual(serialize_split_v1(copy), frame) + + def test_query_auth_nullable_empty_and_large_policy(self): + inner = deserialize_split_v1(_GOLDEN_DATA_SPLIT_V1, self._partition_fields()) + for auth in [None, TableQueryAuthResult(None, None), TableQueryAuthResult([], {}), + TableQueryAuthResult(["η”²" * 30000], {"value": "πŸ™‚" * 20000})]: + with self.subTest(auth=auth): + frame = serialize_split_v1(QueryAuthSplit(inner, auth)) + copy = deserialize_split_v1(frame, self._partition_fields()) + self.assertEqual(serialize_split_v1(copy), frame) + if auth is None: + self.assertIsNone(copy.auth_result) + else: + self.assertEqual(copy.auth_result.filter, auth.filter) + self.assertEqual(copy.auth_result.column_masking, auth.column_masking) + + def test_query_auth_preserves_indexed_selection_and_scores(self): + inner = deserialize_split_v1(_GOLDEN_INDEXED_SPLIT_V1, self._partition_fields()) + frame = serialize_split_v1(QueryAuthSplit(inner, TableQueryAuthResult([], {}))) + copy = deserialize_split_v1(frame, self._partition_fields()) + self.assertIsInstance(copy.split, IndexedSplit) + self.assertEqual(copy.split.row_ranges(), inner.row_ranges()) + self.assertEqual(copy.split.scores(), inner.scores()) + self.assertEqual(serialize_split_v1(copy), frame) + + def test_query_auth_malformed_frame_never_becomes_a_plain_split(self): + inner = serialize_split_v1(deserialize_split_v1(_GOLDEN_DATA_SPLIT_V1, self._partition_fields())) + header = struct.pack('>qii', 0x53504C49545F5631, 1, 5) + valid = header + inner + b'\x01\x00\x00' + self.assertIsInstance(deserialize_split_v1(valid, self._partition_fields()), QueryAuthSplit) + malformed = [valid[:-1], valid + b'\x00', header + inner + b'\x01\x01' + struct.pack('>i', -1), + header + inner + b'\x01\x01' + struct.pack('>i', 2147483647), + header + inner + b'\x01\x01' + struct.pack('>ii', 1, -1) + b'\x00', + header * 33 + inner + b'\x00' * 33] + for frame in malformed: + with self.subTest(frame=frame[-16:]), self.assertRaises(ValueError): + deserialize_split_v1(frame, self._partition_fields()) + def test_deserialize_data_split_v1_golden(self): split = deserialize_split_v1(_GOLDEN_DATA_SPLIT_V1, self._partition_fields())